Add constants
Change-Id: I1a071d77f50566d1d25dbd0ba3a99d1e6ac648b2
diff --git a/example.py b/example.py
index 6dd6f1a..a0a691a 100644
--- a/example.py
+++ b/example.py
@@ -1,18 +1,18 @@
import vision
def run_detector_example():
- detector = vision.Detector('ssd_mobilenet_v2_face_quant_postprocess_edgetpu.tflite')
+ detector = vision.Detector(vision.FACE_DETECTION_MODEL)
for frame in vision.get_frames('Face Detector', size=(640, 480)):
faces = detector.get_objects(frame)
- vision.draw_objects(frame, faces, color=(255, 0, 255), thickness=5)
+ vision.draw_objects(frame, faces)
def run_classifier_example():
- labels = vision.load_labels('imagenet_labels.txt')
- classifier = vision.Classifier('mobilenet_v2_1.0_224_quant_edgetpu.tflite')
+ labels = vision.load_labels(vision.CLASSIFICATION_LABELS)
+ classifier = vision.Classifier(vision.CLASSIFICATION_MODEL)
for frame in vision.get_frames('Object Classifier', size=(640, 480)):
classes = classifier.get_classes(frame)
- vision.draw_classes(frame, classes, labels, color=(255, 0, 255))
+ vision.draw_classes(frame, classes, labels)
if __name__ == '__main__':
- #run_classifier_example()
- run_detector_example()
+ run_classifier_example()
+ #run_detector_example()
diff --git a/vision.py b/vision.py
index d2262ae..c18b160 100644
--- a/vision.py
+++ b/vision.py
@@ -5,6 +5,12 @@
import classify
import detect
+FACE_DETECTION_MODEL = 'ssd_mobilenet_v2_face_quant_postprocess_edgetpu.tflite'
+CLASSIFICATION_MODEL = 'mobilenet_v2_1.0_224_quant_edgetpu.tflite'
+CLASSIFICATION_LABELS = 'imagenet_labels.txt'
+
+CORAL_COLOR = (86, 104, 237)
+
def load_labels(filename, encoding='utf-8'):
with open(filename, 'r', encoding=encoding) as f:
return {index : line.strip() for (index, line) in enumerate(f.readlines())}
@@ -39,12 +45,12 @@
self.interpreter.invoke()
return classify.get_output(self.interpreter, top_k, threshold)
-def draw_objects(frame, objs, color, thickness):
+def draw_objects(frame, objs, color=CORAL_COLOR, thickness=5):
for obj in objs:
bbox = obj.bbox
cv2.rectangle(frame, (bbox.xmin, bbox.ymin), (bbox.xmax, bbox.ymax), color, thickness)
-def draw_classes(frame, classes, labels, color):
+def draw_classes(frame, classes, labels, color=CORAL_COLOR):
for index, score in classes:
label = '%s (%.2f)' % (labels.get(index, 'n/a'), score)
cv2.putText(frame, label, (10, 30), cv2.FONT_HERSHEY_PLAIN, 2.0, color, 2)
