[Asone] Add param to draw motion trails on video.

asone/asone.py

@@ -31,18 +31,19 @@ def get_tracker(self, tracker: int):
                           use_cuda=self.use_cuda)
         return tracker
 
-    def track_video(self, video_path, output_dir='results', save_result=True, display=False):
+    def track_video(self, video_path, output_dir='results', save_result=True, display=False, draw_trails=False):
         output_filename = os.path.basename(video_path)
 
         for (bbox_details, frame_details) in self._start_tracking(video_path,
                                                                     output_filename,
                                                                     output_dir=output_dir,
                                                                     save_result=save_result,
-                                                                    display=display):
+                                                                    display=display,
+                                                                    draw_trails=draw_trails):
             # yeild bbox_details, frame_details to main script
             yield bbox_details, frame_details
 
-    def track_webcam(self, cam_id=0, output_dir='results', save_result=False, display=True):
+    def track_webcam(self, cam_id=0, output_dir='results', save_result=False, display=True, draw_trails=False):
 
         output_filename = 'results.mp4'
 
@@ -51,11 +52,13 @@ def track_webcam(self, cam_id=0, output_dir='results', save_result=False, displa
                                                                     output_dir=output_dir,
                                                                     fps=29,
                                                                     save_result=save_result,
-                                                                    display=display):
+                                                                    display=display,
+                                                                    draw_trails=draw_trails):
             # yeild bbox_details, frame_details to main script
             yield bbox_details, frame_details
 
-    def _start_tracking(self, stream_path, filename,  fps=None, output_dir='results', save_result=True, display=False):
+    def _start_tracking(self, stream_path, filename,  fps=None, output_dir='results',
+                        save_result=True, display=False, draw_trails=False):
 
         cap = cv2.VideoCapture(stream_path)
         width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
@@ -98,7 +101,7 @@ def _start_tracking(self, stream_path, filename,  fps=None, output_dir='results'
                 'frame {}/{} ({:.2f} ms)'.format(frame_id, int(frame_count),
                                                  elapsed_time * 1000), )
 
-            im0 = utils.draw_boxes(im0, bboxes_xyxy, identities=ids)
+            im0 = utils.draw_boxes(im0, bboxes_xyxy, identities=ids, draw_trails=draw_trails)
 
             currTime = time.time()
             fps = 1 / (currTime - prevTime)

asone/utils/draw.py

@@ -3,8 +3,10 @@
 import numpy as np
 from asone.utils import compute_color_for_labels
 from asone.utils import get_names
+from collections import deque
 
 names = get_names()
+data_deque = {}
 
 
 def draw_ui_box(x, img, id, color=None, label=None, line_thickness=None):
@@ -42,9 +44,7 @@ def draw_ui_box(x, img, id, color=None, label=None, line_thickness=None):
 def draw_border(img, pt1, pt2, color, thickness, r, d):
     x1, y1 = pt1
     x2, y2 = pt2
-    # Top left
-    cv2.line(img, (x1 + r, y1), (x1 + r + d, y1), color, thickness)
-    cv2.line(img, (x1, y1 + r), (x1, y1 + r + d), color, thickness)
+    # Top leftfrom collections import deque (x1, y1 + r + d), color, thickness)
     cv2.ellipse(img, (x1 + r, y1 + r), (r, r), 180, 0, 90, color, thickness)
 
     # Top right
@@ -71,10 +71,16 @@ def draw_border(img, pt1, pt2, color, thickness, r, d):
     return img
 
 
-def draw_boxes(img, bbox_xyxy, class_ids=None, identities=None, offset=(0, 0)):
+def draw_boxes(img, bbox_xyxy, class_ids=None, identities=None, draw_trails=False, offset=(0, 0)):
     # cv2.line(img, line2[0], line2[1], (0,200,0), 3)
     height, width, _ = img.shape
+
     # remove tracked point from buffer if object is lost
+    if draw_trails:
+        for key in list(data_deque):
+            if key not in identities:
+                data_deque.pop(key)
+
     for i, box in enumerate(bbox_xyxy):
         x1, y1, x2, y2 = [int(i) for i in box]
         x1 += offset[0]
@@ -85,18 +91,46 @@ def draw_boxes(img, bbox_xyxy, class_ids=None, identities=None, offset=(0, 0)):
         # box_area = (x2-x1) * (y2-y1)
         box_height = (y2-y1)
 
-        # code to find center of bottom edge
-        center = (int((x2+x1) / 2), int((y2+y2)/2))
-
         # get ID of object
         id = int(identities[i]) if identities is not None else None
 
         color = compute_color_for_labels(0)
 
+
         label = None
         if class_ids is not None:
             color = compute_color_for_labels(int(class_ids[i]))
             obj_name = names[int(class_ids[i])]
             label = '%s' % (obj_name)            
         draw_ui_box(box, img, id=id, label=label, color=color, line_thickness=2)
+
+        # Draw trails
+
+        # code to find center of bottom edge
+        center = (int((x2+x1) / 2), int((y2+y2)/2))
+
+        if draw_trails:
+            # create new buffer for new object
+            if id not in data_deque:  
+                data_deque[id] = deque(maxlen= 64)
+
+            data_deque[id].appendleft(center)    
+            drawtrails(data_deque, id, color, img)
+
+
+
     return img
+
+def drawtrails(data_deque, id, color, img):
+    # draw trail
+    for i in range(1, len(data_deque[id])):
+        # check if on buffer value is none
+        if data_deque[id][i - 1] is None or data_deque[id][i] is None:
+            continue
+
+        # generate dynamic thickness of trails
+        thickness = int(np.sqrt(64 / float(i + i)) * 1.5)
+
+        # draw trails
+        cv2.line(img, data_deque[id][i - 1], data_deque[id][i], color, thickness)
+

main.py

@@ -9,12 +9,11 @@ def main(args):
         detector=asone.YOLOX_DARKNET_PYTORCH,
         use_cuda=args.use_cuda
         )
-    for bbox_details, frame_details in dt_obj.track_video(
-                                                            args.video_path,
-                                                            output_dir=args.output_dir,
-                                                            save_result=args.save_result,
-                                                            display=args.display
-                                                            ):
+    for bbox_details, frame_details in dt_obj.track_video(args.video_path,
+                                                          output_dir=args.output_dir,
+                                                          save_result=args.save_result,
+                                                          display=args.display,
+                                                          draw_trails=args.draw_trails):
         # bbox_details = bbox_xyxy, ids, scores, class_ids
         # frame_details = frame, frame_no, fps
 
@@ -31,6 +30,7 @@ def main(args):
     parser.add_argument('--no_display', default=True, action='store_false',
                         dest='display', help='if provided the results will not be displayed on screen')
     parser.add_argument('--output_dir', default='data/results',  help='Path to output directory')
+    parser.add_argument('--draw_trails', default=False,  help='if provided object motion trails will be drawn.')
 
     args = parser.parse_args()
 

setup.py

@@ -9,7 +9,7 @@
 URL = 'https://github.com/axcelerateai/asone'
 DOWNLOAD_URL = URL
 
-VERSION = '0.1.2.dev8'
+VERSION = '0.1.2.dev10'
 
 with open('README.md') as f:
     long_description = f.read()