diff --git a/run_passthru.sh b/run_passthru.sh index 3cf0603..32a14db 100755 --- a/run_passthru.sh +++ b/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 diff --git a/uap_tracker/camera_stream_controller.py b/uap_tracker/camera_stream_controller.py index 1cd116e..958c0ae 100644 --- a/uap_tracker/camera_stream_controller.py +++ b/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 diff --git a/uap_tracker/main.py b/uap_tracker/main.py index a7085d6..7c5f586 100644 --- a/uap_tracker/main.py +++ b/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: diff --git a/uap_tracker/two_by_two_optical_flow_visualiser.py b/uap_tracker/two_by_two_optical_flow_visualiser.py index 4dad198..6d6ffcd 100644 --- a/uap_tracker/two_by_two_optical_flow_visualiser.py +++ b/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 diff --git a/uap_tracker/video_tracker.py b/uap_tracker/video_tracker.py index dc9447f..ad1d721 100644 --- a/uap_tracker/video_tracker.py +++ b/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']