benchmark_tf.py

#!/usr/bin/env python3

import platform
import distro
PLATFORM = platform.system().lower()
GOOGLE = 'edge_tpu'
INTEL = 'ncs2'
NVIDIA = 'jetson_nano'
PI = 'raspberry_pi'
MAC = 'darwin'
IS_LINUX = (PLATFORM == 'linux')

if IS_LINUX:
    PLATFORM = distro.id()
    if PLATFORM == 'debian':
        try:
            with open('/sys/firmware/devicetree/base/model') as f:
                for line in f:
                    line = line.strip()
                    if line.startswith('Raspberry'):
                        PLATFORM = PI
                        print("Running on a Raspberry Pi.")
                        break
        except:
            print("Unknown platform based on Debian.")
            pass
    elif PLATFORM == 'mendel':
        PLATFORM = GOOGLE
        print("Running on a Coral Dev Board.")

try:
    from edgetpu.detection.engine import DetectionEngine
    print("DetectionEngine present.")
except ImportError:
    try:
        from openvino.inference_engine import IENetwork, IEPlugin
        print("OpenVINO present.")
        print("Assuming Movidius hardware.")
        PLATFORM = INTEL
    except ImportError:
        try:
            import tensorflow as tf
            if (tf.test.is_built_with_cuda()):
                import tensorflow.contrib.tensorrt
                print("TensorFlow with GPU support present.")
                print("Assuming Jetson Nano.")
                PLATFORM = NVIDIA
            else:
                print("No GPU support in TensorFlow.")
        except ImportError:
            print("No TensorFlow support found.")

LEGAL_PLATFORMS = NVIDIA, PI, MAC
assert PLATFORM in LEGAL_PLATFORMS, "Don't understand platform %s." % PLATFORM

import sys
import os
import logging as log
import argparse
import subprocess
from timeit import default_timer as timer

import cv2

from PIL import Image
from PIL import ImageFont, ImageDraw

# Function to draw a rectangle with width > 1
def draw_rectangle(draw, coordinates, color, width=1):
    for i in range(width):
        rect_start = (coordinates[0] - i, coordinates[1] - i)
        rect_end = (coordinates[2] + i, coordinates[3] + i)
        draw.rectangle((rect_start, rect_end), outline = color, fill = color)

# 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

def inference_tf(runs, image, model, output, label=None):
   if label:
       labels = ReadLabelFile(label)
   else:
       labels = None
   
   tf_config = tf.compat.v1.ConfigProto()
   tf_config.gpu_options.allow_growth = True
   
   with tf.compat.v1.gfile.FastGFile(model, 'rb') as f:
        graph_def = tf.compat.v1.GraphDef()
        graph_def.ParseFromString(f.read())
        
        
   with tf.compat.v1.Session(config=tf_config) as sess:
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')
        
        img = Image.open(image)
        draw = ImageDraw.Draw(img, 'RGBA')
        helvetica=ImageFont.truetype("./Helvetica.ttf", size=72)
        
        picture = cv2.imread(image)
        initial_h, initial_w, channels = picture.shape
        frame = cv2.resize(picture, (300, 300))
        frame = frame[:, :, [2, 1, 0]]  # BGR2RGB
        frame = frame.reshape(1, frame.shape[0], frame.shape[1], 3)
        
        #  Start synchronous inference and get inference result
        # Run inference.
        print("Running inferencing for ", runs, " times.")
       
        if runs == 1:
           start = timer()
           out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
                           sess.graph.get_tensor_by_name('detection_scores:0'),
                           sess.graph.get_tensor_by_name('detection_boxes:0'),
                           sess.graph.get_tensor_by_name('detection_classes:0')],
                          feed_dict={'image_tensor:0': frame})
           end = timer()
           print('Elapsed time is ', ((end - start)/runs)*1000, 'ms' )
        else:
           start = timer()
           print('Initial run, discarding.')
           out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
                           sess.graph.get_tensor_by_name('detection_scores:0'),
                           sess.graph.get_tensor_by_name('detection_boxes:0'),
                           sess.graph.get_tensor_by_name('detection_classes:0')],
                          feed_dict={'image_tensor:0': frame})
           end = timer()
           print('First run time is ', (end - start)*1000, 'ms')
           start = timer()
           for i in range(runs):
              out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
                              sess.graph.get_tensor_by_name('detection_scores:0'),
                              sess.graph.get_tensor_by_name('detection_boxes:0'),
                              sess.graph.get_tensor_by_name('detection_classes:0')],
                             feed_dict={'image_tensor:0': frame})
           end = timer()
           print('Elapsed time is ', ((end - start)/runs)*1000, 'ms' )
        
        # Visualize detected bounding boxes.
        num_detections = int(out[0][0])
        for i in range(num_detections):
            classId = int(out[3][0][i])
            score = float(out[1][0][i])
            bbox = [float(v) for v in out[2][0][i]]
            if score > 0.5:
                xmin = bbox[1] * initial_w
                ymin = bbox[0] * initial_h
                xmax = bbox[3] * initial_w
                ymax = bbox[2] * initial_h
                if labels:
                   print(labels[classId], 'score = ', score)
                else:
                   print ('score = ', score)
                box = [xmin, ymin, xmax, ymax]
                print( 'box = ', box )
                draw_rectangle(draw, box, (0,128,128,20), width=5)
                if labels:
                   draw.text((box[0] + 20, box[1] + 20), labels[classId], fill=(255,255,255,20), font=helvetica)
        img.save(output)
        print ('Saved to ', output)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', help='Path of the detection model.', required=True)
    parser.add_argument('--label', help='Path of the labels file.')
    parser.add_argument('--input', help='File path of the input image.', required=True)
    parser.add_argument('--output', help='File path of the output image.')
    parser.add_argument('--runs', help='Number of times to run the inference', type=int, default=1)
    args = parser.parse_args()
    
    if ( args.output):
      output_file = args.output
    else:
      output_file = 'out.jpg'
    
    if ( args.label ):
      label_file = args.label
    else:
      label_file = None
    
    result = inference_tf( args.runs, args.input, args.model, output_file, label_file)

if __name__ == '__main__':
  main()