chris/frigate
1
0
Fork 0
forked from Github/frigate
Code Pull Requests Packages Activity

add color map to use different colors for different objects

Browse Source
This commit is contained in:
Blake Blackshear
2019-12-14 16:38:01 -06:00
parent bee99ca6ff
commit 480175d70f
5 changed files with 41 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
frigate/object_detection.py
View File

@@ -4,22 +4,7 @@ import cv2
import threading
import numpy as np
from edgetpu.detection.engine import DetectionEngine
from . util import tonumpyarray
# 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.
PATH_TO_LABELS = '/label_map.pbtext'
# Function to read labels from text files.
def ReadLabelFile(file_path):
with open(file_path, 'r') as f:
lines = f.readlines()
ret = {}
for line in lines:
pair = line.strip().split(maxsplit=1)
ret[int(pair[0])] = pair[1].strip()
return ret
from . util import tonumpyarray, LABELS, PATH_TO_CKPT
class PreppedQueueProcessor(threading.Thread):
def __init__(self, cameras, prepped_frame_queue):
@@ -30,7 +15,7 @@ class PreppedQueueProcessor(threading.Thread):
# Load the edgetpu engine and labels
self.engine = DetectionEngine(PATH_TO_CKPT)
self.labels = ReadLabelFile(PATH_TO_LABELS)
self.labels = LABELS
def run(self):
# process queue...