ML Detector

run_passthru.sh

@@ -1 +1 @@
-PYTHONPATH="${PYTHONPATH}:uap_tracker/" SKY360_VISUALIZER__format=none SKY360_CONTROLLER=camera SKY360_VIDETRACKER__CALCULATE_OPTICAL_FLOW=false SKY360_VIDEOTRACKER__DETECTION_MODE=none SKY360_output_format=video python3 uap_tracker/main.py
+PYTHONPATH="${PYTHONPATH}:uap_tracker/" SKY360_VISUALIZER__format=none SKY360_CONTROLLER=camera SKY360_VIDETRACKER__CALCULATE_OPTICAL_FLOW=false SKY360_VIDEOTRACKER__DETECTION_MODE=none SKY360_VIDEOTRACKER__MASK_PCT=0 SKY360_output_format=video python3 uap_tracker/main.py

uap_tracker/camera_stream_controller.py

@@ -6,8 +6,7 @@
 
 class CameraStreamController():
 
-    def __init__(self, camera, video_tracker, minute_interval=10):
-
+    def __init__(self, camera, video_tracker, minute_interval=1):
         self.camera = camera
         self.video_tracker = video_tracker
         self.minute_interval = minute_interval

uap_tracker/main.py

@@ -5,6 +5,7 @@
 import os
 import getopt
 import sys
+from pathlib import Path
 import cv2
 
 from uap_tracker.event_publisher import EventPublisher
@@ -51,6 +52,7 @@ def _get_visualizer(detection_mode):
     two_by_two_mode_visualizers = {
         'background_subtraction': TwoByTwoVisualiser,
         'optical_flow': TwoByTwoOpticalFlowVisualiser,
+        'mldetector': TwoByTwoOpticalFlowVisualiser,
         'none': None
     }
 
@@ -84,11 +86,13 @@ def _get_detection_mode():
     detection_mode = settings.VideoTracker.get(
         'detection_mode', None)
 
-    detection_modes = ['background_subtraction', 'optical_flow', 'none']
+    detection_modes = ['background_subtraction',
+                       'mldetector',
+                       'none']
 
     if not detection_mode:
         print(
-            f"Please set detection_mode in the config or use the SKY360_DETECTION_MODE env var: {detection_modes}")
+            f"Please set detection_mode in the config or use the SKY360_VIDEOTRACKER__DETECTION_MODE env var: {detection_modes}")
         sys.exit(1)
     else:
         print(f"Detection Mode: {detection_mode}")
@@ -157,20 +161,24 @@ def main(argv):
         sys.exit(1)
 
     try:
-        opts, args = getopt.getopt(argv, "hf:", [])
+        opts, args = getopt.getopt(argv, "hf:d:", [])
     except getopt.GetoptError:
         print(USAGE)
         sys.exit(2)
 
     cmdline_filename = None
+    cmdline_dirname = None
     for opt, arg in opts:
         if opt == '-h':
             print(USAGE)
             sys.exit()
         if opt == '-f':
             cmdline_filename = arg
+        if opt == '-d':
+            cmdline_dirname = arg
 
     print(f"cmdline_filename: {cmdline_filename}")
+    print(f"cmdline_dirname: {cmdline_dirname}")
     print('Settings are ', settings.as_dict())
 
     #cv2.namedWindow("Tracking", cv2.WINDOW_AUTOSIZE)
@@ -187,6 +195,10 @@ def main(argv):
     if cmdline_filename:
         process_file(controller, visualizer, cmdline_filename,
                      output_dir, detection_mode)
+    elif cmdline_dirname:
+        for mp4 in sorted(Path(cmdline_dirname).rglob('video.mp4')):
+            process_file(controller, visualizer, str(mp4),
+                     output_dir, detection_mode)
     else:
         if controller == VideoPlaybackController:
 

uap_tracker/two_by_two_optical_flow_visualiser.py

@@ -10,6 +10,7 @@ def visualise_frame(self, video_tracker):
 
         frame_input = video_tracker.get_image('original')
         optical_flow_frame = video_tracker.get_image('optical_flow')
+        detections_frame = video_tracker.get_image('detections')
         frame_output = video_tracker.get_annotated_image()
         fps = video_tracker.get_fps()
 
@@ -20,7 +21,7 @@ def visualise_frame(self, video_tracker):
             video_tracker, frame_output, self.font_size, self.font_colour, fps)
 
         bottom_left_frame = optical_flow_frame
-        bottom_right_frame = optical_flow_frame
+        bottom_right_frame = detections_frame
 
         return utils.combine_frames_2x2(
             frame_input, frame_output, bottom_left_frame, bottom_right_frame

uap_tracker/video_tracker.py

@@ -7,6 +7,7 @@
 from uap_tracker.tracker import Tracker
 from uap_tracker.background_subtractor_factory import BackgroundSubtractorFactory
 from uap_tracker.dense_optical_flow import DenseOpticalFlow
+from mldetector.model import Model
 from uap_tracker.dense_optical_flow_cuda import DenseOpticalFlowCuda
 
 #
@@ -63,6 +64,8 @@ def __init__(self, detection_mode, events, visualizer, detection_sensitivity=2,
         self.background_subtractor = BackgroundSubtractorFactory.create(
             self.background_subtractor_type, self.detection_sensitivity)
 
+        self.mldetector = None
+
     @property
     def is_tracking(self):
         return len(self.live_trackers) > 0
@@ -116,7 +119,8 @@ def update_trackers(self, tracker_type, bboxes, frame):
         # Mike: We can do the tracker updates in parallel
         for i in range(tracker_count):
             tracker = self.live_trackers[i]
-            threads[i] = Thread(target=self.update_tracker_task, args=(tracker, frame, results, i))
+            threads[i] = Thread(target=self.update_tracker_task,
+                                args=(tracker, frame, results, i))
             threads[i].start()
 
         # Mike: We got to wait for the threads to join before we proceed
@@ -223,7 +227,7 @@ def process_frame(self, frame, frame_count, fps):
     def process_frame_cuda(self, frame, frame_count, fps):
 
         with Stopwatch(mask='CUDA Frame '+str(frame_count)+': Took {s:0.4f} seconds to process', quiet=True):
-            # print(f" fps:{int(fps)}", end='\r')
+            print(f" fps:{int(fps)}", end='\r')
             self.fps = fps
             self.frame_count = frame_count
             frame_w = frame.shape[0]
@@ -237,9 +241,10 @@ def process_frame_cuda(self, frame, frame_count, fps):
             gpu_frame = cv2.cuda_GpuMat()
             gpu_frame.upload(frame)
 
-            scale, scaled_width, scaled_height = utils.calc_image_scale(frame_w, frame_h, self.normalised_w_h[0], self.normalised_w_h[1])
-            if scale:
-                gpu_frame = cv2.cuda.resize(gpu_frame, (scaled_width, scaled_height))
+            if self.resize_frame:
+                scale, scaled_width, scaled_height = utils.calc_image_scale(frame_w, frame_h, self.normalised_w_h[0], self.normalised_w_h[1])
+                if scale:
+                    gpu_frame = cv2.cuda.resize(gpu_frame, (scaled_width, scaled_height))
 
             # Mike: Able to offload to CUDA
             gpu_frame_grey = cv2.cuda.cvtColor(gpu_frame, cv2.COLOR_BGR2GRAY)
@@ -274,6 +279,39 @@ def process_frame_cuda(self, frame, frame_count, fps):
                     optical_flow_cuda_thread.start()
                     worker_threads.append(optical_flow_cuda_thread)
 
+            elif self.detection_mode == 'mldetector':
+                if not self.mldetector:
+                    num_classes = 2
+                    input_channels = 6
+                    self.mldetector = Model(input_channels, num_classes)
+                    checkpoint = self.mldetector.get_checkpoint("checkpoint.pth")
+                    self.mldetector.resume(checkpoint["model"])
+                self.frames['optical_flow'] = self.optical_flow(frame_grey)
+                labels = self.mldetector.detect(self.frames['original'], self.frames['optical_flow'])
+                #print(labels)
+                np_bboxes = labels[0]['boxes'].detach().cpu().numpy()
+                np_scores = labels[0]['scores'].detach().cpu().numpy()
+                bboxes = []
+                for np_box, np_score in zip(np_bboxes, np_scores):
+                    if np_score > 0.10:
+                        x1 = int(np_box[0])
+                        y1 = int(np_box[1])
+                        x2 = int(np_box[2])
+                        y2 = int(np_box[3])
+
+                        bboxes.append([
+                            x1,
+                            y1,
+                            x2-x1,
+                            y2-y1])
+                        #print(f"bboxes: {bboxes}")
+                keypoints = []
+                detections_frame = frame.copy()
+                for box in bboxes:
+                    utils.add_bbox_to_image(
+                        box, detections_frame, 0, 1, (0, 0, 255))
+                    self.frames['detections'] = detections_frame
+
             else:
                 bboxes = []
                 keypoints = []
@@ -349,10 +387,12 @@ def get_annotated_image(self, active_trackers_only=True):
             annotated_frame = self.frames['original'].copy()
             if active_trackers_only:
                 for tracker in self.active_trackers():
-                    utils.add_bbox_to_image(tracker.get_bbox(), annotated_frame, tracker.id, 1, tracker.bbox_color())
+                    utils.add_bbox_to_image(
+                        tracker.get_bbox(), annotated_frame, tracker.id, 1, tracker.bbox_color())
             else:
                 for tracker in self.live_trackers:
-                    utils.add_bbox_to_image(tracker.get_bbox(), annotated_frame, tracker.id, 1, tracker.bbox_color())
+                    utils.add_bbox_to_image(
+                        tracker.get_bbox(), annotated_frame, tracker.id, 1, tracker.bbox_color())
             self.frames['annotated_image'] = annotated_frame
 
         return self.frames['annotated_image']