mirror of
https://github.com/blakeblackshear/frigate.git
synced 2025-08-08 13:51:01 +02:00
added ability to save detection samples for offline analisys and custom model traiing
This commit is contained in:
ivelin
parent
6b6519ecfe
commit
e26baeb782
@ -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")
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user