Update app.py

app.py

@@ -1,120 +1,73 @@
 import os
 import tempfile
-from typing import Dict, Tuple
-
 import cv2
-import gradio as gr
 import numpy as np
+import gradio as gr
 import supervision as sv
 from ultralytics import YOLO
 
-
-MODEL_PATH = "yolo11n.pt"
-model = YOLO(MODEL_PATH)
-
+model = YOLO("yolo11n.pt")
 CLASS_NAMES_DICT = model.model.names
-SELECTED_CLASS_NAMES = ["car", "bus", "truck", "motorcycle"]
+
+SELECTED_CLASS_NAMES = ['car', 'bus', 'truck', 'motorcycle']
 SELECTED_CLASS_IDS = [
-    {value: key for key, value in CLASS_NAMES_DICT.items()}[class_name]
-    for class_name in SELECTED_CLASS_NAMES
+    {value: key for key, value in CLASS_NAMES_DICT.items()}[name]
+    for name in SELECTED_CLASS_NAMES
 ]
 
+box_annotator = sv.BoxAnnotator(thickness=4)
+label_annotator = sv.LabelAnnotator(text_thickness=2, text_scale=1.5, text_color=sv.Color.BLACK)
+trace_annotator = sv.TraceAnnotator(thickness=4, trace_length=50)
+line_zone_annotator = sv.LineZoneAnnotator(
+    thickness=4, text_thickness=2, text_scale=2,
+    display_in_count=False, display_out_count=False
+)
 
-def process_video(
-    video_path: str,
-    frame_stride: int = 2,
-    max_seconds: int = 60,
-    line_orientation: str = "horizontal",
-) -> Tuple[str, Dict[str, int]]:
-    """
-    Xử lý video: phát hiện + tracking + đếm phương tiện.
-    Trả về: đường dẫn video đã annotate và dict số lượng theo class.
-    """
-    video_info = sv.VideoInfo.from_video_path(video_path)
-    fps = video_info.fps
-    max_frames_for_detection = int(fps * max_seconds) if max_seconds is not None else None
 
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise RuntimeError("Không thể mở video đầu vào.")
+def process_video(video_path, orientation="Horizontal"):
+    """Process video: detect objects and return annotated video."""
+    if video_path is None:
+        return None
+
+    if isinstance(video_path, dict):
+        video_path = video_path.get("path", video_path)
 
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
+    video_info = sv.VideoInfo.from_video_path(video_path)
+    w, h = video_info.width, video_info.height
 
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    tmp_dir = tempfile.mkdtemp()
-    out_path = os.path.join(tmp_dir, "result.mp4")
-    writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
+    is_horizontal = orientation.lower().startswith("n")
+    if is_horizontal:
+        line_pos = int(h * 0.5)
+        line_start = sv.Point(int(w * 0.005), line_pos)
+        line_end = sv.Point(int(w * 0.995), line_pos)
+    else:
+        line_pos = int(w * 0.5)
+        line_start = sv.Point(line_pos, int(h * 0.005))
+        line_end = sv.Point(line_pos, int(h * 0.995))
 
+    line_zone = sv.LineZone(start=line_start, end=line_end)
     byte_tracker = sv.ByteTrack(
         track_activation_threshold=0.25,
         lost_track_buffer=30,
         minimum_matching_threshold=0.8,
-        frame_rate=fps,
-        minimum_consecutive_frames=3,
+        frame_rate=video_info.fps or 30,
+        minimum_consecutive_frames=3
     )
+    byte_tracker.reset()
 
-    if line_orientation == "vertical":
-        line_pos = int(width * 0.5)
-        line_start = sv.Point(line_pos, int(height * 0.05))
-        line_end = sv.Point(line_pos, int(height * 0.95))
-        axis = "x"
-    else:
-        line_pos = int(height * 0.5)
-        line_start = sv.Point(int(width * 0.05), line_pos)
-        line_end = sv.Point(int(width * 0.95), line_pos)
-        axis = "y"
-
-    line_zone = sv.LineZone(start=line_start, end=line_end)
-
-    previous_positions: Dict[int, float] = {}
-    class_counts: Dict[str, int] = {name: 0 for name in SELECTED_CLASS_NAMES}
+    previous_positions = {}
+    class_counts = {name: 0 for name in SELECTED_CLASS_NAMES}
     crossed_ids = set()
 
-    box_annotator = sv.BoxAnnotator(thickness=4)
-    label_annotator = sv.LabelAnnotator(
-        text_thickness=2, text_scale=1.0, text_color=sv.Color.BLACK
-    )
-    trace_annotator = sv.TraceAnnotator(thickness=4, trace_length=50)
-    line_zone_annotator = sv.LineZoneAnnotator(
-        thickness=4,
-        color=sv.Color.RED,
-        text_thickness=2,
-        text_scale=2,
-        display_in_count=False,
-        display_out_count=False,
-    )
-
-    frame_idx = 0
-    last_detections = None
+    def callback(frame: np.ndarray, index: int) -> np.ndarray:
+        nonlocal previous_positions, class_counts, crossed_ids
 
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
+        results = model(frame, verbose=False)[0]
+        detections = sv.Detections.from_ultralytics(results)
+        detections = detections[np.isin(detections.class_id, SELECTED_CLASS_IDS)]
+        detections = byte_tracker.update_with_detections(detections)
 
-        annotated_frame = frame.copy()
-
-        can_run_detection = True
-        if max_frames_for_detection is not None and frame_idx >= max_frames_for_detection:
-            can_run_detection = False
-
-        if can_run_detection and frame_idx % frame_stride == 0:
-            results = model(
-                frame,
-                imgsz=640,
-                device="cpu",
-                verbose=False,
-                classes=SELECTED_CLASS_IDS,
-            )[0]
-            detections = sv.Detections.from_ultralytics(results)
-            detections = byte_tracker.update_with_detections(detections)
-            last_detections = detections
-        else:
-            detections = last_detections
-
-        if detections is not None and detections.tracker_id is not None:
+        if detections.tracker_id is not None:
             xyxy = detections.xyxy
             for i in range(len(detections)):
                 tid = int(detections.tracker_id[i])
@@ -123,169 +76,80 @@ def process_video(
                 cx = (xyxy[i, 0] + xyxy[i, 2]) / 2
                 cy = (xyxy[i, 1] + xyxy[i, 3]) / 2
 
-                cur_pos = cy if axis == "y" else cx
-
+                curr_coord = cy if is_horizontal else cx
                 if tid in previous_positions:
-                    prev_pos = previous_positions[tid]
-                    if (
-                        prev_pos < line_pos
-                        and cur_pos > line_pos
-                        and (tid, "pos") not in crossed_ids
-                    ):
-                        crossed_ids.add((tid, "pos"))
+                    prev_coord = previous_positions[tid]
+                    if prev_coord < line_pos and curr_coord > line_pos and (tid, 'pos') not in crossed_ids:
+                        crossed_ids.add((tid, 'pos'))
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                    elif (
-                        prev_pos > line_pos
-                        and cur_pos < line_pos
-                        and (tid, "neg") not in crossed_ids
-                    ):
-                        crossed_ids.add((tid, "neg"))
+                    elif prev_coord > line_pos and curr_coord < line_pos and (tid, 'neg') not in crossed_ids:
+                        crossed_ids.add((tid, 'neg'))
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                previous_positions[tid] = cur_pos
-
-            labels = [
-                f"#{tracker_id} {CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
-                for confidence, class_id, tracker_id in zip(
-                    detections.confidence, detections.class_id, detections.tracker_id
-                )
-            ]
+                previous_positions[tid] = curr_coord
 
-            annotated_frame = trace_annotator.annotate(
-                scene=annotated_frame, detections=detections
-            )
-            annotated_frame = box_annotator.annotate(
-                scene=annotated_frame, detections=detections
-            )
-            annotated_frame = label_annotator.annotate(
-                scene=annotated_frame, detections=detections, labels=labels
+        labels = [
+            f"#{tid} {CLASS_NAMES_DICT[cid]} {conf:0.2f}"
+            for conf, cid, tid in zip(
+                detections.confidence, detections.class_id, detections.tracker_id
             )
+        ]
 
-            line_zone.trigger(detections)
-            annotated_frame = line_zone_annotator.annotate(
-                annotated_frame, line_counter=line_zone
-            )
+        annotator_frame = frame.copy()
+        annotator_frame = trace_annotator.annotate(scene=annotator_frame, detections=detections)
+        annotator_frame = box_annotator.annotate(scene=annotator_frame, detections=detections)
+        annotator_frame = label_annotator.annotate(scene=annotator_frame, detections=detections, labels=labels)
+
+        line_zone.trigger(detections)
+        annotator_frame = line_zone_annotator.annotate(annotator_frame, line_counter=line_zone)
 
-        h, w, _ = annotated_frame.shape
+        fh, fw, _ = annotator_frame.shape
         box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
-        x0, y0 = w - box_w - 20, 20
+        x0, y0 = fw - box_w - 20, 20
 
-        overlay = annotated_frame.copy()
+        overlay = annotator_frame.copy()
         cv2.rectangle(overlay, (x0, y0), (x0 + box_w, y0 + box_h), (0, 0, 0), -1)
-        annotated_frame = cv2.addWeighted(overlay, 0.6, annotated_frame, 0.4, 0)
+        annotator_frame = cv2.addWeighted(overlay, 0.6, annotator_frame, 0.4, 0)
 
         total = sum(class_counts.values())
-        cv2.putText(
-            annotated_frame,
-            f"Total: {total}",
-            (x0 + 10, y0 + 30),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.8,
-            (0, 255, 0),
-            2,
-        )
+        cv2.putText(annotator_frame, f'Total: {total}', (x0 + 10, y0 + 30),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
         for i, cls_name in enumerate(SELECTED_CLASS_NAMES):
             cnt = class_counts.get(cls_name, 0)
-            cv2.putText(
-                annotated_frame,
-                f"{cls_name.capitalize()}: {cnt}",
-                (x0 + 10, y0 + 60 + i * 28),
-                cv2.FONT_HERSHEY_SIMPLEX,
-                0.7,
-                (255, 255, 255),
-                2,
-            )
-
-        writer.write(annotated_frame)
-        frame_idx += 1
-
-    cap.release()
-    writer.release()
-
-    return out_path, class_counts
-
-
-def gradio_infer(video, mode: str = "Cân bằng", line_type: str = "Ngang"):
-    if video is None:
-        return None, "Vui lòng upload 1 video."
-
-    if isinstance(video, dict):
-        video_path = video.get("name") or video.get("data")
-    else:
-        video_path = video
-
-    if not video_path or not os.path.exists(video_path):
-        return None, "Không tìm thấy file video."
+            cv2.putText(annotator_frame, f'{cls_name.capitalize()}: {cnt}',
+                        (x0 + 10, y0 + 60 + i * 28),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
 
-    if mode == "Nhanh":
-        frame_stride = 4
-        max_seconds = 30
-    elif mode == "Chính xác":
-        frame_stride = 1
-        max_seconds = 90
-    else:  # Cân bằng
-        frame_stride = 2
-        max_seconds = 60
+        return annotator_frame
 
-    if line_type == "Dọc":
-        line_orientation = "vertical"
-    else:
-        line_orientation = "horizontal"
-
-    out_path, counts = process_video(
-        video_path=video_path,
-        frame_stride=frame_stride,
-        max_seconds=max_seconds,
-        line_orientation=line_orientation,
+    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    sv.process_video(
+        source_path=video_path,
+        target_path=output_path,
+        callback=callback
     )
+    return output_path
 
-    total = sum(counts.values())
-    lines = [f"Tổng số phương tiện: {total}"]
-    for cls_name in SELECTED_CLASS_NAMES:
-        lines.append(f"- {cls_name}: {counts.get(cls_name, 0)}")
-    summary = "\n".join(lines)
-
-    return out_path, summary
 
-
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # Nhận diện phương tiện trong video (YOLO + ByteTrack)
-
-        Upload 1 video.
-        """
-    )
+with gr.Blocks(title="Object Detection", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚗 Nhận dạng phương tiện (YOLOv8 + ByteTrack)")
+    gr.Markdown("Upload video, hệ thống sẽ nhận dạng phương tiện trong video.")
 
     with gr.Row():
-        with gr.Column():
-            video_input = gr.Video(label="Video đầu vào", sources=["upload"])
-            mode_input = gr.Radio(
-                ["Nhanh", "Cân bằng", "Chính xác"],
-                value="Cân bằng",
-                label="Chế độ xử lý",
-                info="Nhanh: nhanh hơn, ít chính xác hơn. Chính xác: chậm hơn.",
-            )
-            line_type_input = gr.Radio(
-                ["Ngang", "Dọc"],
-                value="Ngang",
-                label="Hướng đường đếm",
-                info="Ngang: đường ngang ở giữa khung hình. Dọc: đường dọc ở giữa khung hình.",
-            )
-            run_btn = gr.Button("Bắt đầu đếm")
-
-        with gr.Column():
-            video_output = gr.Video(label="Video đã xử lý")
-            text_output = gr.Textbox(
-                label="Kết quả đếm", lines=6, interactive=False
-            )
+        video_input = gr.Video(label="Video input")
+        video_output = gr.Video(label="Video output")
 
-    run_btn.click(
-        fn=gradio_infer,
-        inputs=[video_input, mode_input, line_type_input],
-        outputs=[video_output, text_output],
+    orientation_input = gr.Radio(
+        choices=["Horizontal", "Vertical"],
+        value="Horizontal",
+        label="Line orientation"
     )
 
+    btn = gr.Button("▶️ Process Video")
+    btn.click(fn=process_video, inputs=[video_input, orientation_input], outputs=video_output)
+
+    gr.Markdown("""
+    ### Author: Trần Hải Nam - 223332840
+    """)
 
 if __name__ == "__main__":
     demo.launch()
-