This program write does live object detection and write output to a csv file.

object_csv.py

# Copyright 2019 Google LLC
#
# Changes made by Nam Vu 02/25/2020
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This program does live object detection and write output to a csv file.
#
# Example run:
#
# 1) Download model and label:
# Model and Label File taken from here:
# https://github.com/google-coral/edgetpu/tree/master/test_data
# $ wget https://github.com/google-coral/edgetpu/raw/master/test_data/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite
# $ curl https://raw.githubusercontent.com/google-coral/edgetpu/master/test_data/coco_labels.txt > coco_labels.txt
#
# 2) Run:
# python3 object_csv.py --model ./mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite --label ./coco_labels.txt
#
# NOT PRODUCTION READY, JUST AN EXAMPLE DEMO

import argparse
import platform
import subprocess
from edgetpu.detection.engine import DetectionEngine
from edgetpu.utils import dataset_utils
from datetime import datetime
import cv2
import csv
import os
import numpy as np
from PIL import Image
from PIL import ImageDraw

def do_detection(engine, img, k=3):
  # Run inference.
  ans = engine.detect_with_image(
      img,
      threshold=0.05,
      keep_aspect_ratio=False,
      relative_coord=False,
      top_k=k)
  return ans

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--model',
      help='Path of the detection model, it must be a SSD model with postprocessing operator.',
      required=True)
  parser.add_argument('--label', help='Path of the labels file.')
  parser.add_argument('--top_k', help='Top number to results to return.')
  args = parser.parse_args()

  # Initialize engine.
  engine = DetectionEngine(args.model)
  labels = dataset_utils.read_label_file(args.label) if args.label else None
  top_k = args.top_k if args.top_k else 3
  
  # make csv file
  filename = "output.csv"
  writer = 0 
  if os.path.exists(filename):
    result = open("output.csv", "a")
    writer = csv.writer(result)
  else:
    result = open("output.csv", "a")  
    writer = csv.writer(result)
    row = ['object', 'score', 'timestamp']
    writer.writerow(row)

  # capturing videos with cv2
  cap = cv2.VideoCapture(0)
  # looping through all frames
  frame_num = 0
  while True:
    print('<<<----------------------------------------->>>')
    print('Frame number: ', frame_num)
    ret, cv2img = cap.read()
    cv2img = cv2.cvtColor(cv2img, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(cv2img)
    draw = ImageDraw.Draw(img)
    ans = do_detection(engine, img)
    # Save result.
    if ans:
      for obj in ans:
        print('-----------------------------------------')
        if labels:
          print(labels[obj.label_id])
          print('score = ', obj.score)
          box = obj.bounding_box.flatten().tolist()
          print('box=', box)
          draw.rectangle(box, outline='red')
          t = datetime.now()
          row = [str(labels[obj.label_id]), str(obj.score), str(t)]
          writer.writerow(row)
    else:
      print('No object detected!')
    new_cv2img = np.array(img.convert('RGB')) 
    new_cv2img = new_cv2img[:, :, ::-1].copy() # RGB to BGR
    cv2.imshow('Object Detection', new_cv2img)
    if cv2.waitKey(1) & 0xFF == ord('q'):
      break
    frame_num += 1
    print('\n')

  cap.release()
  cv2.destroyAllWindows()

if __name__ == '__main__':
  main()