Procházet zdrojové kódy

computer vision examples

Andrey Koryagin před 6 měsíci
rodič
revize
0573c8cd6f

binární
plate_detection/license_plate_detector.pt


+ 63 - 0
plate_detection/plate_detection.py

@@ -0,0 +1,63 @@
+from ultralytics import YOLO
+import cv2
+import easyocr
+
+# Load a model
+model = YOLO('license_plate_detector.pt')  # load an official model
+
+# Open the input video
+cap = cv2.VideoCapture('trafik.mp4')
+
+reader = easyocr.Reader(['ch_sim','en']) # this needs to run only once to load the model into memory
+
+bad_boys = ['GXISOGJ','GXIS OGJ', 'GXI5OGJ', 'GXI5 OGJ', 'EYGINBG']
+
+if not cap.isOpened():
+    raise Exception("Error: Could not open video.")
+
+
+while cap.isOpened():
+    # Read a frame from the input video
+    success, frame = cap.read()
+
+    if success:
+        results = model.predict(source=frame, conf=0.30, verbose=False)
+
+        for result in results:
+            box = result.boxes.xyxy.cpu().numpy().astype(int)
+            if (len(box) > 0 ) :
+                x1 = box[0][0]
+                y1 = box[0][1]
+                x2 = box[0][2]
+                y2 = box[0][3]
+
+                new_img = frame[y1:y2, x1:x2]
+                #new_img = cv2.cvtColor(new_img, cv2.COLOR_BGR2GRAY)
+                #new_img = cv2.convertScaleAbs(new_img, alpha=1.5, beta=30.0)
+
+                result = reader.readtext(new_img)
+
+                plate_text = ''
+                for (bbox, text, prob) in result:
+                    plate_text += text
+                print (plate_text)
+
+                if plate_text in bad_boys:
+                    color = (0, 0, 255)
+                    frame = cv2.putText(frame, 'BAD BOY', (x1, y1-50), cv2.FONT_HERSHEY_SIMPLEX, 3.0, color, 10)
+                else:
+                    color = (0, 255, 0)
+
+                frame = cv2.rectangle(frame, (x1, y1), (x2, y2), color, 4)
+                frame = cv2.putText(frame, plate_text, (x1, y1-3), cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)
+
+                #cv2.imshow("Plate", new_img)
+
+        frame = cv2.resize(frame, (0, 0), fx = 0.3, fy = 0.3)
+        cv2.imshow("License plate detector", frame)
+        key = cv2.waitKey(1)
+        if key & 0xff == ord('q'):
+            break
+
+cap.release()
+cv2.destroyAllWindows()