plate_detection.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. from ultralytics import YOLO
  2. import cv2
  3. import easyocr
  4. # Load a model
  5. model = YOLO('license_plate_detector.pt') # load an official model
  6. # Open the input video
  7. cap = cv2.VideoCapture('trafik.mp4')
  8. reader = easyocr.Reader(['ch_sim','en']) # this needs to run only once to load the model into memory
  9. bad_boys = ['GXISOGJ','GXIS OGJ', 'GXI5OGJ', 'GXI5 OGJ', 'EYGINBG']
  10. if not cap.isOpened():
  11. raise Exception("Error: Could not open video.")
  12. while cap.isOpened():
  13. # Read a frame from the input video
  14. success, frame = cap.read()
  15. if success:
  16. results = model.predict(source=frame, conf=0.30, verbose=False)
  17. for result in results:
  18. box = result.boxes.xyxy.cpu().numpy().astype(int)
  19. if (len(box) > 0 ) :
  20. x1 = box[0][0]
  21. y1 = box[0][1]
  22. x2 = box[0][2]
  23. y2 = box[0][3]
  24. new_img = frame[y1:y2, x1:x2]
  25. #new_img = cv2.cvtColor(new_img, cv2.COLOR_BGR2GRAY)
  26. #new_img = cv2.convertScaleAbs(new_img, alpha=1.5, beta=30.0)
  27. result = reader.readtext(new_img)
  28. plate_text = ''
  29. for (bbox, text, prob) in result:
  30. plate_text += text
  31. print (plate_text)
  32. if plate_text in bad_boys:
  33. color = (0, 0, 255)
  34. frame = cv2.putText(frame, 'BAD BOY', (x1, y1-50), cv2.FONT_HERSHEY_SIMPLEX, 3.0, color, 10)
  35. else:
  36. color = (0, 255, 0)
  37. frame = cv2.rectangle(frame, (x1, y1), (x2, y2), color, 4)
  38. frame = cv2.putText(frame, plate_text, (x1, y1-3), cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)
  39. #cv2.imshow("Plate", new_img)
  40. frame = cv2.resize(frame, (0, 0), fx = 0.3, fy = 0.3)
  41. cv2.imshow("License plate detector", frame)
  42. key = cv2.waitKey(1)
  43. if key & 0xff == ord('q'):
  44. break
  45. cap.release()
  46. cv2.destroyAllWindows()