diff --git a/frigate/object_detection.py b/frigate/object_detection.py index 6c598a321..500bc341c 100755 --- a/frigate/object_detection.py +++ b/frigate/object_detection.py @@ -6,11 +6,16 @@ import threading import numpy as np from edgetpu.detection.engine import DetectionEngine from . util import tonumpyarray +import os +import json -# Path to frozen detection graph. This is the actual model that is used for the object detection. +# Default path to frozen detection graph. This is the actual model that is used for the object detection. PATH_TO_CKPT = '/frozen_inference_graph.pb' -# List of the strings that is used to add correct label for each box. +# Default path to list of the strings that is used to add correct label for each box. PATH_TO_LABELS = '/label_map.pbtext' +# Default data_dir used to save detections for later use such as re-labeling and training +DATA_DIR = '/data/' + # Function to read labels from text files. def ReadLabelFile(file_path): @@ -35,15 +40,20 @@ class PreppedQueueProcessor(threading.Thread): if tf_args != None: tf_model = tf_args.get('model', None) tf_labels = tf_args.get('labels', None) + self.data_dir = tf_args.get('data_dir', DATA_DIR) else: tf_model = PATH_TO_CKPT tf_labels = PATH_TO_LABELS + self.data_dir = DATA_DIR print("Loading Tensorflow model: " + str(tf_model)) self.engine = DetectionEngine(tf_model) if tf_labels: self.labels = ReadLabelFile(tf_labels) else: self.labels = None + yesterday = datetime.datetime.now() - datetime.timedelta(days=1) + self.last_saved_det_frame_time = yesterday + self.last_saved_non_det_frame_time = yesterday def run(self): # process queue... @@ -54,35 +64,107 @@ class PreppedQueueProcessor(threading.Thread): # objects = self.engine.DetectWithInputTensor(frame['frame'], threshold=frame['region_threshold'], top_k=3) # print(self.engine.get_inference_time()) - objects = self.engine.DetectWithImage( - frame['img'], - threshold=frame['region_threshold'], - keep_aspect_ratio=True, - relative_coord=True, - top_k=3) + try: + # print("Calling TF Coral to detect objects") + detected_objects = self.engine.DetectWithImage( + frame['img'], + threshold=frame['region_threshold'], + keep_aspect_ratio=True, + relative_coord=True, + top_k=3) + #print("TF Coral detected {} objects withing {} ms". + # format(len(objects), + # self.engine.get_inference_time())) - # parse and pass detected objects back to the camera - parsed_objects = [] - for obj in objects: - print("Detected an object: \n\n") - print("Detected object label index: " + str(obj.label_id) + "\n\n") - box = obj.bounding_box.flatten().tolist() - if self.labels: - obj_name = str(self.labels[obj.label_id]) + # parse and pass detected objects back to the outbound camera stream + self.__pass_detections_to_output_cam(frame, detected_objects) + + # save detections to file if enabled + self.__save_sample(frame, detected_objects) + except Exception as e: + print("Error while calling TF Coral: \n" + str(e)) + + # Parse and pass detected objects back to the outbound camera stream + def __pass_detections_to_output_cam(self, frame, detected_objects): + parsed_objects = [] + for obj in detected_objects: + # print("Detected an object: \n\n") + # print("Detected object label index: " + str(obj.label_id) + "\n\n") + box = obj.bounding_box.flatten().tolist() + if self.labels: + obj_name = str(self.labels[obj.label_id]) + else: + obj_name = " " # no labels, just a yes/no type of detection + detection = { + 'frame_time': frame['frame_time'], + 'name': obj_name, + 'score': float(obj.score), + 'xmin': int((box[0] * frame['region_size']) + frame['region_x_offset']), + 'ymin': int((box[1] * frame['region_size']) + frame['region_y_offset']), + 'xmax': int((box[2] * frame['region_size']) + frame['region_x_offset']), + 'ymax': int((box[3] * frame['region_size']) + frame['region_y_offset']) + } + # print(str(detection) + "\n") + parsed_objects.append(detection) + self.cameras[frame['camera_name']].add_objects(parsed_objects) + + # Save detections to file for offline processing + def __save_sample(self, frame, detected_objects): + try: + if frame['save_samples']: + now = datetime.datetime.now() + dir = self.data_dir + frame['camera_name'] + os.makedirs(dir, exist_ok=True) + # pretty up the timestamp + file_prefix = now.strftime("%Y%m%d-%H%M%S") + if detected_objects: + # default 2 seconds between saved detection samples + save_det_interval = frame['save_samples'].get('detection_interval', 2) + # print("Detection interval for saved samples : {}".format(save_det_interval)) + # check if enough time has passed since the last saved detection sample + if (now - self.last_saved_det_frame_time).total_seconds() >= save_det_interval: + dir += "/detections/" + os.makedirs(dir, exist_ok=True) + det_file_path = dir + file_prefix + ".json" + detection_json = [] + for obj in detected_objects: + box = obj.bounding_box.flatten().tolist() + if self.labels: + obj_name = str(self.labels[obj.label_id]) + else: + obj_name = " " # no labels, just a yes/no type of detection + next_detection = { + 'time': now.isoformat(' '), + 'label_id': obj.label_id, + 'name': obj_name, + 'score': float(obj.score), + 'xmin': int((box[0] * frame['region_size'])), + 'ymin': int((box[1] * frame['region_size'])), + 'xmax': int((box[2] * frame['region_size'])), + 'ymax': int((box[3] * frame['region_size'])) + } + detection_json.append(next_detection) + with open(det_file_path, 'w', encoding='utf-8') as f: + json.dump(detection_json, f, ensure_ascii=False, indent=4) + + self.last_saved_det_frame_time = now + img_file_path = dir + file_prefix + ".jpg" + img = frame['img'] + img.save(img_file_path, "JPEG") else: - obj_name = " " # no labels, just a yes/no type of detection - detection = { - 'frame_time': frame['frame_time'], - 'name': obj_name, - 'score': float(obj.score), - 'xmin': int((box[0] * frame['region_size']) + frame['region_x_offset']), - 'ymin': int((box[1] * frame['region_size']) + frame['region_y_offset']), - 'xmax': int((box[2] * frame['region_size']) + frame['region_x_offset']), - 'ymax': int((box[3] * frame['region_size']) + frame['region_y_offset']) - } - print(str(detection) + "\n") - parsed_objects.append(detection) - self.cameras[frame['camera_name']].add_objects(parsed_objects) + # default 300 seconds (5 min) between saved non detection samples + save_non_det_interval = frame['save_samples'].get('detection_interval', 300) + # print("Detection interval for non saved samples : {}".format(save_non_det_interval)) + # check if enough time has passed since the last saved non-detection sample + if (now - self.last_saved_non_det_frame_time).total_seconds() >= save_non_det_interval: + dir += "/non_detections/" + os.makedirs(dir, exist_ok=True) + self.last_saved_non_det_frame_time = now + img_file_path = dir + file_prefix + ".jpg" + img = frame['img'] + img.save(img_file_path, "JPEG") + except Exception as e: + print("Error while saving an image region sample: \n" + str(e)) # should this be a region class? @@ -90,6 +172,7 @@ class FramePrepper(threading.Thread): def __init__(self, camera_name, shared_frame, frame_time, frame_ready, frame_lock, region_size, region_x_offset, region_y_offset, region_threshold, + region_save_samples, prepped_frame_queue): threading.Thread.__init__(self) @@ -102,6 +185,7 @@ class FramePrepper(threading.Thread): self.region_x_offset = region_x_offset self.region_y_offset = region_y_offset self.region_threshold = region_threshold + self.region_save_samples = region_save_samples self.prepped_frame_queue = prepped_frame_queue def run(self): @@ -136,7 +220,8 @@ class FramePrepper(threading.Thread): 'region_size': self.region_size, 'region_threshold': self.region_threshold, 'region_x_offset': self.region_x_offset, - 'region_y_offset': self.region_y_offset + 'region_y_offset': self.region_y_offset, + 'save_samples': self.region_save_samples }) else: print("queue full. moving on") diff --git a/frigate/video.py b/frigate/video.py index 8b0f8e444..5fe45919d 100755 --- a/frigate/video.py +++ b/frigate/video.py @@ -163,6 +163,7 @@ class Camera: self.frame_ready, self.frame_lock, region['size'], region['x_offset'], region['y_offset'], region['threshold'], + region['save_samples'], prepped_frame_queue )) @@ -320,7 +321,7 @@ class Camera: draw_box_with_label(frame, obj['xmin'], obj['ymin'], obj['xmax'], obj['ymax'], label) for region in self.regions: - color = (255,255,255) + color = (255,255,255) # white bounding box for regions where we look for objects cv2.rectangle(frame, (region['x_offset'], region['y_offset']), (region['x_offset']+region['size'], region['y_offset']+region['size']), color, 2)