Tensorflow Object Detection Webcam

# Hai sa importam cele de trebuinta

import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image


from utils import label_map_util

from utils import visualization_utils as vis_util

# Pregatim modelul

'''
Orice model exportat folosind 'export_inference_graph.py'
poate fi incarcat aici prin schimbarea 'PATH_TO_CKPT'
in asa fel incat noua destinatie sa corespunda
noului fisier .pb

Eu am folosit aici modelul 'SSD with Mobilenet'
'''

# Stabilim ce model sa downloadam.
MODEL_NAME = 'ssd_mobilenet_v1_coco_11_06_2017'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'

#Path-ul catre modelul folosit pentru identificarea obiectelor
#In cazul nostru, frozen detection graph

PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'

#Path-ul catre lista de denumiri ale obiectelor

PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')

NUM_CLASSES = 90


# Acush sa download modelu'

if not os.path.exists(MODEL_NAME + '/frozen_inference_graph.pb'):
    print ('Downloading the model')
    opener = urllib.request.URLopener()
    opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
    tar_file = tarfile.open(MODEL_FILE)
    for file in tar_file.getmembers():
      file_name = os.path.basename(file.name)
      if 'frozen_inference_graph.pb' in file_name:
        tar_file.extract(file, os.getcwd())
    print ('Download complete')
else:
    print ('Model already exists')

# Sa bagam in memorie un model Tensorflow (frozen)

detection_graph = tf.Graph()
with detection_graph.as_default():
  od_graph_def = tf.GraphDef()
  with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
    serialized_graph = fid.read()
    od_graph_def.ParseFromString(serialized_graph)
    tf.import_graph_def(od_graph_def, name='')


# Ce e aia Label map?
# Label map este o lista de indici. Cand reteaua noastra neuronala
# face o predictie, sa zicem 5, asta inseamna ca
# a gasit un avion.

label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)

# Sa initializam webcam-ul...

import cv2
cap = cv2.VideoCapture(0)

# ... si sa punem tensorflow la treaba!

with detection_graph.as_default():
  with tf.Session(graph=detection_graph) as sess:
   ret = True
   while (ret):
      ret,image_np = cap.read()
      # Trebuie sa expandam putin dimensiunile, din moment ce modelul
      # se asteapta ca imaginile sa aiba o forma: [1, None, None, 3]
      image_np_expanded = np.expand_dims(image_np, axis=0)
      image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
      # Fiecare contur (box) reprezinta o parte a imaginii
      # unde a fost detectat un obiect cunoscut
      boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
      # Fiecare scor reprezinta nivelul de incredere,
      # ca ce scrie e tot aia cu ce se vede
      # Adica identificarea este corecta intr-o proportie de x%
      # Scorul apare langa obiect impreuna cu denumirea obiectului (eticheta).
      scores = detection_graph.get_tensor_by_name('detection_scores:0')
      classes = detection_graph.get_tensor_by_name('detection_classes:0')
      num_detections = detection_graph.get_tensor_by_name('num_detections:0')
      # Aici se produce detectia propriu-zisa
      (boxes, scores, classes, num_detections) = sess.run(
          [boxes, scores, classes, num_detections],
          feed_dict={image_tensor: image_np_expanded})
      # Aici se gaseste vizualizarea
      vis_util.visualize_boxes_and_labels_on_image_array(
          image_np,
          np.squeeze(boxes),
          np.squeeze(classes).astype(np.int32),
          np.squeeze(scores),
          category_index,
          use_normalized_coordinates=True,
          line_thickness=8)
      cv2.imshow('image',cv2.resize(image_np,(1280,960)))
      if cv2.waitKey(25) & 0xFF == ord('q'):
          cv2.destroyAllWindows()
          cap.release()
          break

# Ura si la gara!