Update app.py

app.py

@@ -8,6 +8,7 @@ from ultralytics import YOLO
 
 DEFAULT_MAX_FRAME_SIZE = 640
 DEFAULT_DETECT_EVERY_N_FRAMES = 2
+DEFAULT_CONF_THRESHOLD = 0.3
 
 model = YOLO("yolov8n.pt")
 CLASS_NAMES_DICT = model.model.names
@@ -27,6 +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"
 ):
     if video_path is None:
         return None
@@ -49,22 +51,38 @@ def process_video(
     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
+
     def callback(frame: np.ndarray, index: int) -> np.ndarray:
         nonlocal previous_positions, class_counts, crossed_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, max_size / max(fh_orig, fw_orig))
+            scale = min(1.0, effective_max_size / max(fh_orig, fw_orig))
             if scale < 1.0:
                 frame_infer = cv2.resize(
                     frame, (int(fw_orig * scale), int(fh_orig * scale))
@@ -135,7 +153,11 @@ 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)
 
         if detections.tracker_id is not None:
@@ -175,7 +197,8 @@ def process_video(
         ]
 
         annotator_frame = frame_infer.copy()
-        annotator_frame = trace_annotator.annotate(scene=annotator_frame, detections=detections)
+        if enable_trace:
+            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
@@ -278,26 +301,39 @@ 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)",
+                label="Kích thước tối đa (px) (chỉ dùng khi 'Tuỳ chỉnh')",
             )
             detect_every_n = gr.Slider(
                 minimum=1,
                 maximum=5,
                 value=DEFAULT_DETECT_EVERY_N_FRAMES,
                 step=1,
-                label="Detect mỗi N frame (1 = mọi frame)",
+                label="Detect mỗi N frame (chỉ dùng khi 'Tuỳ chỉnh')",
             )
 
     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],
+        inputs=[
+            video_input,
+            use_resize,
+            max_frame_size,
+            detect_every_n,
+            line_orientation,
+            performance_mode,
+        ],
         outputs=video_output,
     )