main.py 1004 B

1234567891011121314151617181920212223242526272829303132333435
  1. import os
  2. import shutil
  3. import re
  4. from yolov5 import detect
  5. import easyocr
  6. def main():
  7. directory = "plates"
  8. detect.run(
  9. source=directory,
  10. save_crop=True,
  11. weights="yolo-licence-plate.pt",
  12. conf_thres=0.25,
  13. iou_thres=0.45,
  14. agnostic_nms=False,
  15. max_det=1,
  16. project="predictions",
  17. exist_ok=True
  18. )
  19. reader = easyocr.Reader(["es"])
  20. crops_directory = "predictions/exp/crops/license_plate"
  21. for image in os.listdir(crops_directory):
  22. filepath = os.path.join(crops_directory, image)
  23. if os.path.isfile(filepath) and image.endswith((".jpg", ".jpeg", ".png")):
  24. image_path = os.path.join(crops_directory, image)
  25. print(f"Processing {image_path}...")
  26. result = reader.readtext(image_path, detail=0)
  27. clean_plate = re.sub(r"[^A-Z0-9]", "", result[0].upper())
  28. print(clean_plate)
  29. shutil.rmtree("predictions/exp")
  30. if __name__ == "__main__":
  31. main()