Update app.py

app.py

@@ -8,7 +8,7 @@ from ultralytics import YOLO
 
 DEFAULT_MAX_FRAME_SIZE = 640
 DEFAULT_DETECT_EVERY_N_FRAMES = 2
-DEFAULT_CONF_THRESHOLD = 0.3
+DEFAULT_ZONE_MARGIN = 0.10  # Độ dày vùng đếm quanh line (tỉ lệ chiều cao/rộng)
 
 model = YOLO("yolov8n.pt")
 CLASS_NAMES_DICT = model.model.names
@@ -28,7 +28,7 @@ def process_video(
     max_frame_size: int = DEFAULT_MAX_FRAME_SIZE,
     detect_every_n: int = DEFAULT_DETECT_EVERY_N_FRAMES,
     line_orientation: str = "Ngang",
-    performance_mode: str = "Cân bằng",  # "Nhanh" | "Cân bằng" | "Đẹp" | "Tuỳ chỉnh"
+    zone_margin: float = DEFAULT_ZONE_MARGIN,
 ):
     if video_path is None:
         return None
@@ -47,42 +47,25 @@ def process_video(
     )
     byte_tracker.reset()
 
-    previous_positions = {}
     class_counts = {name: 0 for name in SELECTED_CLASS_NAMES}
-    crossed_ids = set()
-
-    # Cấu hình hiệu năng theo chế độ
-    if performance_mode == "Nhanh":
-        effective_max_size = 480
-        detect_every = 4
-        conf_threshold = 0.5
-        enable_trace = False
-    elif performance_mode == "Đẹp":
-        effective_max_size = 800
-        detect_every = 1
-        conf_threshold = 0.3
-        enable_trace = True
-    elif performance_mode == "Cân bằng":
-        effective_max_size = DEFAULT_MAX_FRAME_SIZE
-        detect_every = DEFAULT_DETECT_EVERY_N_FRAMES
-        conf_threshold = 0.4
-        enable_trace = True
-    else:  # Tuỳ chỉnh
-        effective_max_size = (
-            int(max_frame_size) if max_frame_size and max_frame_size > 0 else DEFAULT_MAX_FRAME_SIZE
-        )
-        detect_every = (
-            int(detect_every_n) if detect_every_n and detect_every_n >= 1 else 1
-        )
-        conf_threshold = DEFAULT_CONF_THRESHOLD
-        enable_trace = True
+    counted_ids = set()
 
     def callback(frame: np.ndarray, index: int) -> np.ndarray:
-        nonlocal previous_positions, class_counts, crossed_ids
+        nonlocal class_counts, counted_ids
+
+        if max_frame_size is None or max_frame_size <= 0:
+            max_size = DEFAULT_MAX_FRAME_SIZE
+        else:
+            max_size = int(max_frame_size)
+
+        if detect_every_n is None or detect_every_n < 1:
+            detect_every = 1
+        else:
+            detect_every = int(detect_every_n)
 
         fh_orig, fw_orig = frame.shape[:2]
         if use_resize:
-            scale = min(1.0, effective_max_size / max(fh_orig, fw_orig))
+            scale = min(1.0, max_size / max(fh_orig, fw_orig))
             if scale < 1.0:
                 frame_infer = cv2.resize(
                     frame, (int(fw_orig * scale), int(fh_orig * scale))
@@ -93,6 +76,7 @@ def process_video(
             frame_infer = frame
 
         fh, fw = frame_infer.shape[:2]
+        # Vị trí line theo hướng người dùng chọn
         if line_orientation == "Dọc":
             line_pos = int(fw * 0.5)
             is_horizontal = False
@@ -100,24 +84,89 @@ def process_video(
             line_pos = int(fh * 0.5)
             is_horizontal = True
 
+        # Vùng đếm (zone) quanh line, dày theo tỉ lệ zone_margin
+        if zone_margin is None or zone_margin <= 0:
+            zm_ratio = DEFAULT_ZONE_MARGIN
+        else:
+            zm_ratio = max(0.01, min(0.5, float(zone_margin)))
+
+        if is_horizontal:
+            z_half = int(fh * zm_ratio)
+            z_top = max(0, line_pos - z_half)
+            z_bot = min(fh - 1, line_pos + z_half)
+        else:
+            z_half = int(fw * zm_ratio)
+            z_left = max(0, line_pos - z_half)
+            z_right = min(fw - 1, line_pos + z_half)
+
         if detect_every > 1 and index % detect_every != 0:
             annotator_frame = frame_infer.copy()
+            # Vẽ vùng đếm (zone) + line trung tâm (không cập nhật đếm để tiết kiệm CPU)
+            overlay_zone = annotator_frame.copy()
+            if is_horizontal:
+                cv2.rectangle(
+                    overlay_zone,
+                    (0, z_top),
+                    (fw, z_bot),
+                    (0, 0, 200),
+                    -1,
+                )
+            else:
+                cv2.rectangle(
+                    overlay_zone,
+                    (z_left, 0),
+                    (z_right, fh),
+                    (0, 0, 200),
+                    -1,
+                )
+            annotator_frame = cv2.addWeighted(
+                overlay_zone, 0.18, annotator_frame, 0.82, 0
+            )
 
+            thickness_base = max(2, int(2 * (max(fw, fh) / 1920)))
             if is_horizontal:
                 cv2.line(
                     annotator_frame,
-                    (int(fw * 0.05), line_pos),
-                    (int(fw * 0.95), line_pos),
-                    (0, 255, 255),
-                    2,
+                    (0, z_top),
+                    (fw, z_top),
+                    (0, 100, 255),
+                    thickness_base,
+                )
+                cv2.line(
+                    annotator_frame,
+                    (0, z_bot),
+                    (fw, z_bot),
+                    (0, 100, 255),
+                    thickness_base,
+                )
+                cv2.line(
+                    annotator_frame,
+                    (0, line_pos),
+                    (fw, line_pos),
+                    (0, 0, 255),
+                    max(3, thickness_base + 1),
                 )
             else:
                 cv2.line(
                     annotator_frame,
-                    (line_pos, int(fh * 0.05)),
-                    (line_pos, int(fh * 0.95)),
-                    (0, 255, 255),
-                    2,
+                    (z_left, 0),
+                    (z_left, fh),
+                    (0, 100, 255),
+                    thickness_base,
+                )
+                cv2.line(
+                    annotator_frame,
+                    (z_right, 0),
+                    (z_right, fh),
+                    (0, 100, 255),
+                    thickness_base,
+                )
+                cv2.line(
+                    annotator_frame,
+                    (line_pos, 0),
+                    (line_pos, fh),
+                    (0, 0, 255),
+                    max(3, thickness_base + 1),
                 )
 
             box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
@@ -153,13 +202,10 @@ def process_video(
 
         results = model(frame_infer, verbose=False)[0]
         detections = sv.Detections.from_ultralytics(results)
-        # Lọc theo class quan tâm
         detections = detections[np.isin(detections.class_id, SELECTED_CLASS_IDS)]
-        # Lọc thêm theo confidence để giảm số lượng box
-        if len(detections) > 0 and hasattr(detections, "confidence"):
-            detections = detections[detections.confidence >= conf_threshold]
         detections = byte_tracker.update_with_detections(detections)
 
+        # Đếm theo loại khi đi vào vùng đếm (zone) quanh line
         if detections.tracker_id is not None:
             xyxy = detections.xyxy
             for i in range(len(detections)):
@@ -169,25 +215,14 @@ def process_video(
                 cx = (xyxy[i, 0] + xyxy[i, 2]) / 2
                 cy = (xyxy[i, 1] + xyxy[i, 3]) / 2
 
-                curr_coord = cy if is_horizontal else cx
-
-                if tid in previous_positions:
-                    prev_coord = previous_positions[tid]
-                    if (
-                        prev_coord < line_pos
-                        and curr_coord > line_pos
-                        and (tid, "out") not in crossed_ids
-                    ):
-                        crossed_ids.add((tid, "out"))
+                # Đếm một lần khi ID lần đầu đi vào vùng đếm
+                if cls_name in SELECTED_CLASS_NAMES and tid not in counted_ids:
+                    if is_horizontal and z_top <= cy <= z_bot:
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                    elif (
-                        prev_coord > line_pos
-                        and curr_coord < line_pos
-                        and (tid, "in") not in crossed_ids
-                    ):
-                        crossed_ids.add((tid, "in"))
+                        counted_ids.add(tid)
+                    elif (not is_horizontal) and z_left <= cx <= z_right:
                         class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
-                previous_positions[tid] = curr_coord
+                        counted_ids.add(tid)
 
         labels = [
             f"#{tid} {CLASS_NAMES_DICT[cid]} {conf:0.2f}"
@@ -197,28 +232,78 @@ def process_video(
         ]
 
         annotator_frame = frame_infer.copy()
-        if enable_trace:
-            annotator_frame = trace_annotator.annotate(scene=annotator_frame, detections=detections)
+        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
         )
 
+        # Vẽ vùng đếm (zone) + line trung tâm
+        overlay_zone = annotator_frame.copy()
+        if is_horizontal:
+            cv2.rectangle(
+                overlay_zone,
+                (0, z_top),
+                (fw, z_bot),
+                (0, 0, 200),
+                -1,
+            )
+        else:
+            cv2.rectangle(
+                overlay_zone,
+                (z_left, 0),
+                (z_right, fh),
+                (0, 0, 200),
+                -1,
+            )
+        annotator_frame = cv2.addWeighted(
+            overlay_zone, 0.18, annotator_frame, 0.82, 0
+        )
+
+        thickness_base = max(2, int(2 * (max(fw, fh) / 1920)))
         if is_horizontal:
             cv2.line(
                 annotator_frame,
-                (int(fw * 0.05), line_pos),
-                (int(fw * 0.95), line_pos),
-                (0, 255, 255),
-                2,
+                (0, z_top),
+                (fw, z_top),
+                (0, 100, 255),
+                thickness_base,
+            )
+            cv2.line(
+                annotator_frame,
+                (0, z_bot),
+                (fw, z_bot),
+                (0, 100, 255),
+                thickness_base,
+            )
+            cv2.line(
+                annotator_frame,
+                (0, line_pos),
+                (fw, line_pos),
+                (0, 0, 255),
+                max(3, thickness_base + 1),
             )
         else:
             cv2.line(
                 annotator_frame,
-                (line_pos, int(fh * 0.05)),
-                (line_pos, int(fh * 0.95)),
-                (0, 255, 255),
-                2,
+                (z_left, 0),
+                (z_left, fh),
+                (0, 100, 255),
+                thickness_base,
+            )
+            cv2.line(
+                annotator_frame,
+                (z_right, 0),
+                (z_right, fh),
+                (0, 100, 255),
+                thickness_base,
+            )
+            cv2.line(
+                annotator_frame,
+                (line_pos, 0),
+                (line_pos, fh),
+                (0, 0, 255),
+                max(3, thickness_base + 1),
             )
 
         box_w, box_h = 280, 50 + len(SELECTED_CLASS_NAMES) * 28
@@ -301,39 +386,33 @@ with gr.Blocks(title="Nhận dạng phương tiện giao thông", theme=gr.theme
                 value="Ngang",
                 label="Hướng phương tiện di chuyển",
             )
-        with gr.Row():
-            performance_mode = gr.Radio(
-                choices=["Nhanh", "Cân bằng", "Đẹp", "Tuỳ chỉnh"],
-                value="Cân bằng",
-                label="Chế độ hiệu năng",
-            )
         with gr.Row():
             max_frame_size = gr.Slider(
                 minimum=320,
                 maximum=1280,
                 value=DEFAULT_MAX_FRAME_SIZE,
                 step=64,
-                label="Kích thước tối đa (px) (chỉ dùng khi 'Tuỳ chỉnh')",
+                label="Kích thước tối đa (px)",
             )
             detect_every_n = gr.Slider(
                 minimum=1,
                 maximum=5,
                 value=DEFAULT_DETECT_EVERY_N_FRAMES,
                 step=1,
-                label="Detect mỗi N frame (chỉ dùng khi 'Tuỳ chỉnh')",
+                label="Detect mỗi N frame (1 = mọi frame)",
+            )
+            zone_margin = gr.Slider(
+                minimum=0.02,
+                maximum=0.30,
+                value=DEFAULT_ZONE_MARGIN,
+                step=0.01,
+                label="Độ dày vùng đếm quanh line",
             )
 
     btn = gr.Button("▶️ Xử lý video", variant="primary")
     btn.click(
         fn=process_video,
-        inputs=[
-            video_input,
-            use_resize,
-            max_frame_size,
-            detect_every_n,
-            line_orientation,
-            performance_mode,
-        ],
+        inputs=[video_input, use_resize, max_frame_size, detect_every_n, line_orientation, zone_margin],
         outputs=video_output,
     )