fix

detectron/demo/main.py

@@ -98,9 +98,11 @@ def create_ui():
             with gr.Row():
                 with gr.Column(scale=2):
                     with gr.Row(elem_id='audio-container'):
-                        with gr.Group():
+                        with gr.Column(scale=2):
                             components["yolo_dist_a"] = gr.Number(value=0,label="A",visible=False)
+                        with gr.Column(scale=2):
                             components["yolo_dist_b"] = gr.Number(value=0,label="B",visible=False)
+                        with gr.Group():
                             components["yolo_image_input"] = gr.Image(type="pil",elem_id='image-input',label='输入')
                             components["yolo_video_input"] = gr.Video(label='输入',visible=False,interactive=True)
 
@@ -307,8 +309,9 @@ def do_yolo_algo_type_chage(value):
         components["yolo_video_output"] = gr.PlayableVideo(label='输出',visible=True)
         components["yolo_image_input"] = gr.Image(type="pil",elem_id='image-input',label='输入',visible=False)
         components["yolo_image_output"] = gr.Image(type="pil",elem_id='image-output',label='输出',interactive=False,visible=True)
-        components["yolo_dist_a"] = gr.Number(value=0,label="A",visible=True)
-        components["yolo_dist_b"] = gr.Number(value=0,label="B",visible=True)
+        if "距离" == value.strip():
+            components["yolo_dist_a"] = gr.Number(value=0,label="A",visible=True,info="跟踪id-A")
+            components["yolo_dist_b"] = gr.Number(value=0,label="B",visible=True,info="跟踪id-B")
     else:
         components["yolo_image_input"] = gr.Image(type="pil",elem_id='image-input',label='输入',visible=True)
         components["yolo_image_output"] = gr.Image(type="pil",elem_id='image-output',label='输出',interactive=False,visible=True)

detectron/demo/yolo_predictor.py

@@ -195,15 +195,14 @@ def do_gym(frame):
 
         yield from self._video_processor(video_path,do_gym)
 
-    def heatmap(self,video_path):
+    def heatmap(self,video_path,classes_for_heatmap = None,region_points=None):
         heatmap_obj = solutions.Heatmap(
             colormap=cv2.COLORMAP_PARULA,
             view_img=False,
             shape="circle",
             classes_names=self.model.names,
+            count_reg_pts=region_points,
         )
-
-        classes_for_heatmap = [0, 2]
 
         def do_draw(frame):
             try:
@@ -228,7 +227,7 @@ def do_draw(frame):
 
         yield from self._video_processor(video_path,do_draw)
 
-    def vision_eye(self,video_path):
+    def vision_eye(self,video_path,center_point=None):
         import math
 
         pixel_per_meter = 10
@@ -237,8 +236,9 @@ def vision_eye(self,video_path):
         def vision_distance(frame):
             annotator = Annotator(frame, line_width=2)
 
-            height = frame.shape[0]
-            center_point = (0, height)
+            if center_point is None:
+                height = frame.shape[0]
+                center_point = (0, height)
 
             results = self.model.track(frame, persist=True)
             boxes = results[0].boxes.xyxy.cpu()
@@ -302,7 +302,7 @@ def queue(frame):
         yield from self._video_processor(video_path,queue)
 
     def _post_processor(self, output):
-        print("-------yolo------------\n", output)
+        # print("-------yolo------------\n", output)
         pil_images = []
 
         result: Dict[str, Instances] = {