video function

app.py

@@ -3,17 +3,93 @@ import cv2
 from detector.utils import detect_plates, detect_chars
 
 
-def predict(img):
-    plates = detect_plates(img)
+def fn_image(foto):
+    plates = detect_plates(foto)
+    plates_text = []
     if len(plates) > 0:
         for plate in plates:
             p1, p2, crop = plate
             if len(crop) > 0:
-                cv2.rectangle(img, p1, p2, (0, 0, 255), 2)
+                cv2.rectangle(foto, p1, p2, (0, 0, 255), 2)
                 text, crop = detect_chars(crop)
-                cv2.putText(img, text, p1, cv2.FONT_HERSHEY_SIMPLEX, 4, (0, 255, 0), 5)
-    return img
+                cv2.putText(foto, text, p1, cv2.FONT_HERSHEY_SIMPLEX, 4, (0, 255, 0), 5)
+                plates_text.append(text)
+    return foto, plates_text
 
 
-iface = gr.Interface(fn=predict, inputs="image", outputs="image")
-iface.launch()
+def fn_video(video, initial_time, duration):
+    plates_text = []
+    cap = cv2.VideoCapture(video)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    image_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+    final_video = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'VP90'), fps, image_size)
+    num_frames = 0
+    min_frame = int(initial_time * fps)
+    max_frame = int((initial_time + duration) * fps)
+    while cap.isOpened():
+        try:
+            ret, frame = cap.read()
+            if not ret:
+                break
+        except Exception as e:
+            print(e)
+            continue
+        if num_frames < min_frame:
+            num_frames += 1
+            continue
+        plates = detect_plates(frame)
+        for plate in plates:
+            p1, p2, crop = plate
+            if len(crop) > 0:
+                cv2.rectangle(frame, p1, p2, (0, 0, 255), 2)
+                text, crop = detect_chars(crop)
+                cv2.putText(frame, text, p1, cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 5)
+                plates_text.append(text)
+        final_video.write(frame)
+        num_frames += 1
+        if num_frames == max_frame:
+            break
+    cap.release()
+    final_video.release()
+    return 'output.mp4', plates_text
+
+
+image_interface = gr.Interface(
+    fn=fn_image,
+    inputs="image",
+    outputs=["image", "text"],
+    title="Buscar números de placa en una imagen",
+    allow_flagging=False,
+    allow_screenshot=False,
+)
+
+video_interface = gr.Interface(
+    fn=fn_video,
+    inputs=[
+        gr.Video(type="file", label="Video"),
+        gr.Slider(0, 600, value=0, label="Tiempo inicial en segundos", step=1),
+        gr.Slider(0, 10, value=4, label="Duración en segundos", step=1),
+    ],
+    outputs=["video", "text"],
+    title="Buscar números de placa en un video",
+    allow_flagging=False,
+    allow_screenshot=False,
+)
+
+webcam_interface = gr.Interface(
+    fn_image,
+    inputs=[
+        gr.Image(source='webcam', streaming=True),
+    ],
+    outputs=gr.Image(type="file"),
+    live=True,
+    title="Buscar placa con la cámara",
+    allow_flagging=False,
+    allow_screenshot=False,
+)
+
+if __name__ == "__main__":
+    gr.TabbedInterface(
+        [image_interface, video_interface],
+        ["Fotos", "Videos"],
+    ).launch()

detector/utils.py

@@ -7,13 +7,6 @@ BASE_DIR = os.path.abspath(os.getcwd())
 
 model_plates = torch.hub.load('ultralytics/yolov5', 'custom',
                               path=os.path.join(BASE_DIR, 'detector', 'static', 'plates.pt'))
-# model.conf = 0.60  # NMS confidence threshold
-# model.iou = 0.60  # NMS IoU threshold
-# model.agnostic = False  # NMS class-agnostic
-# model.multi_label = False  # NMS multiple labels per box
-# model.classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
-# model.max_det = 1  # maximum number of detections per image
-# model.amp = False  # Automatic Mixed Precision (AMP) inference
 
 model_chars = torch.hub.load('ultralytics/yolov5', 'custom',
                              path=os.path.join(BASE_DIR, 'detector', 'static', 'chars.pt'))
@@ -49,7 +42,7 @@ def detect_plates(img):
 
 
 def detect_chars(img):
-    img = cv2.resize(img, (640, 480))
+    img = cv2.resize(img, (640, 320))
     detect = model_chars(img)
     records = detect.pandas().xyxy[0].to_dict(orient='records')
     yolo = np.squeeze(detect.render())
@@ -57,48 +50,9 @@ def detect_chars(img):
     if records:
         records = sorted(records, key=lambda d: d['xmin'])
         text = ''.join([i.get('name') for i in records])
-    return text, img
+    return text, yolo
 
 
 def save_plates(img):
     detect = model_plates(img)
     detect.crop(save=True)
-
-
-# def yolo_detections_to_norfair_detections(yolo_detections, track_points="centroid"):
-#     """convert detections_as_xywh to norfair detections"""
-#     norfair_detections = []
-#
-#     if track_points == "centroid":
-#         detections_as_xywh = yolo_detections.xywh[0]
-#         for detection_as_xywh in detections_as_xywh:
-#             centroid = np.array(
-#                 [detection_as_xywh[0].item(), detection_as_xywh[1].item()]
-#             )
-#             scores = np.array([detection_as_xywh[4].item()])
-#             norfair_detections.append(
-#                 Detection(
-#                     points=centroid,
-#                     scores=scores,
-#                     label=int(detection_as_xywh[-1].item()),
-#                 )
-#             )
-#     elif track_points == "bbox":
-#         detections_as_xyxy = yolo_detections.xyxy[0]
-#         for detection_as_xyxy in detections_as_xyxy:
-#             bbox = np.array(
-#                 [
-#                     [detection_as_xyxy[0].item(), detection_as_xyxy[1].item()],
-#                     [detection_as_xyxy[2].item(), detection_as_xyxy[3].item()],
-#                 ]
-#             )
-#             scores = np.array(
-#                 [detection_as_xyxy[4].item(), detection_as_xyxy[4].item()]
-#             )
-#             norfair_detections.append(
-#                 Detection(
-#                     points=bbox, scores=scores, label=int(detection_as_xyxy[-1].item())
-#                 )
-#             )
-#
-#     return norfair_detections