import cv2
import vision

def run_detector_example():
  detector = vision.make_detector('ssd_mobilenet_v2_face_quant_postprocess_edgetpu.tflite')
  for frame in vision.Camera('Face Detector', size=(640, 480)):
    faces = detector(frame)
    for face in faces:
      bbox = face.bbox
      cv2.rectangle(frame, (bbox.xmin, bbox.ymin), (bbox.xmax, bbox.ymax), (255, 0, 255), 5)

def run_classifier_example():
  labels = vision.load_labels('imagenet_labels.txt')
  classifier = vision.make_classifier('mobilenet_v2_1.0_224_quant_edgetpu.tflite')
  for frame in vision.Camera('Object Classifier', size=(640, 480)):
    classes = classifier(frame)
    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, (255, 0, 255), 2)

if __name__ == '__main__':
  #run_classifier_example()
  run_detector_example()