Merge pull request #9750 from dkurt:feature_dnn_tf_text_graph
Vadim Pisarevsky
8 years ago
commit
046045239c
10 changed files with 538 additions and 85 deletions
@ -0,0 +1,131 @@ |
||||
# Script to evaluate MobileNet-SSD object detection model trained in TensorFlow |
||||
# using both TensorFlow and OpenCV. Example: |
||||
# |
||||
# python mobilenet_ssd_accuracy.py \ |
||||
# --weights=frozen_inference_graph.pb \ |
||||
# --prototxt=ssd_mobilenet_v1_coco.pbtxt \ |
||||
# --images=val2017 \ |
||||
# --annotations=annotations/instances_val2017.json |
||||
# |
||||
# Tested on COCO 2017 object detection dataset, http://cocodataset.org/#download |
||||
import os |
||||
import cv2 as cv |
||||
import json |
||||
import argparse |
||||
|
||||
parser = argparse.ArgumentParser( |
||||
description='Evaluate MobileNet-SSD model using both TensorFlow and OpenCV. ' |
||||
'COCO evaluation framework is required: http://cocodataset.org') |
||||
parser.add_argument('--weights', required=True, |
||||
help='Path to frozen_inference_graph.pb of MobileNet-SSD model. ' |
||||
'Download it at https://github.com/tensorflow/models/tree/master/research/object_detection') |
||||
parser.add_argument('--prototxt', help='Path to ssd_mobilenet_v1_coco.pbtxt from opencv_extra.', required=True) |
||||
parser.add_argument('--images', help='Path to COCO validation images directory.', required=True) |
||||
parser.add_argument('--annotations', help='Path to COCO annotations file.', required=True) |
||||
args = parser.parse_args() |
||||
|
||||
### Get OpenCV predictions ##################################################### |
||||
net = cv.dnn.readNetFromTensorflow(args.weights, args.prototxt) |
||||
|
||||
detections = [] |
||||
for imgName in os.listdir(args.images): |
||||
inp = cv.imread(os.path.join(args.images, imgName)) |
||||
rows = inp.shape[0] |
||||
cols = inp.shape[1] |
||||
inp = cv.resize(inp, (300, 300)) |
||||
|
||||
net.setInput(cv.dnn.blobFromImage(inp, 1.0/127.5, (300, 300), (127.5, 127.5, 127.5), True)) |
||||
out = net.forward() |
||||
|
||||
for i in range(out.shape[2]): |
||||
score = float(out[0, 0, i, 2]) |
||||
# Confidence threshold is in prototxt. |
||||
classId = int(out[0, 0, i, 1]) |
||||
|
||||
x = out[0, 0, i, 3] * cols |
||||
y = out[0, 0, i, 4] * rows |
||||
w = out[0, 0, i, 5] * cols - x |
||||
h = out[0, 0, i, 6] * rows - y |
||||
detections.append({ |
||||
"image_id": int(imgName.rstrip('0')[:imgName.rfind('.')]), |
||||
"category_id": classId, |
||||
"bbox": [x, y, w, h], |
||||
"score": score |
||||
}) |
||||
|
||||
with open('cv_result.json', 'wt') as f: |
||||
json.dump(detections, f) |
||||
|
||||
### Get TensorFlow predictions ################################################# |
||||
import tensorflow as tf |
||||
|
||||
with tf.gfile.FastGFile(args.weights) as f: |
||||
# Load the model |
||||
graph_def = tf.GraphDef() |
||||
graph_def.ParseFromString(f.read()) |
||||
|
||||
with tf.Session() as sess: |
||||
# Restore session |
||||
sess.graph.as_default() |
||||
tf.import_graph_def(graph_def, name='') |
||||
|
||||
detections = [] |
||||
for imgName in os.listdir(args.images): |
||||
inp = cv.imread(os.path.join(args.images, imgName)) |
||||
rows = inp.shape[0] |
||||
cols = inp.shape[1] |
||||
inp = cv.resize(inp, (300, 300)) |
||||
inp = inp[:, :, [2, 1, 0]] # BGR2RGB |
||||
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': inp.reshape(1, inp.shape[0], inp.shape[1], 3)}) |
||||
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.01: |
||||
x = bbox[1] * cols |
||||
y = bbox[0] * rows |
||||
w = bbox[3] * cols - x |
||||
h = bbox[2] * rows - y |
||||
detections.append({ |
||||
"image_id": int(imgName.rstrip('0')[:imgName.rfind('.')]), |
||||
"category_id": classId, |
||||
"bbox": [x, y, w, h], |
||||
"score": score |
||||
}) |
||||
|
||||
with open('tf_result.json', 'wt') as f: |
||||
json.dump(detections, f) |
||||
|
||||
### Evaluation part ############################################################ |
||||
|
||||
# %matplotlib inline |
||||
import matplotlib.pyplot as plt |
||||
from pycocotools.coco import COCO |
||||
from pycocotools.cocoeval import COCOeval |
||||
import numpy as np |
||||
import skimage.io as io |
||||
import pylab |
||||
pylab.rcParams['figure.figsize'] = (10.0, 8.0) |
||||
|
||||
annType = ['segm','bbox','keypoints'] |
||||
annType = annType[1] #specify type here |
||||
prefix = 'person_keypoints' if annType=='keypoints' else 'instances' |
||||
print 'Running demo for *%s* results.'%(annType) |
||||
|
||||
#initialize COCO ground truth api |
||||
cocoGt=COCO(args.annotations) |
||||
|
||||
#initialize COCO detections api |
||||
for resFile in ['tf_result.json', 'cv_result.json']: |
||||
print resFile |
||||
cocoDt=cocoGt.loadRes(resFile) |
||||
|
||||
cocoEval = COCOeval(cocoGt,cocoDt,annType) |
||||
cocoEval.evaluate() |
||||
cocoEval.accumulate() |
||||
cocoEval.summarize() |
||||
@ -0,0 +1,62 @@ |
||||
# This file is part of OpenCV project. |
||||
# It is subject to the license terms in the LICENSE file found in the top-level directory |
||||
# of this distribution and at http://opencv.org/license.html. |
||||
# |
||||
# Copyright (C) 2017, Intel Corporation, all rights reserved. |
||||
# Third party copyrights are property of their respective owners. |
||||
import tensorflow as tf |
||||
import struct |
||||
import argparse |
||||
import numpy as np |
||||
|
||||
parser = argparse.ArgumentParser(description='Convert weights of a frozen TensorFlow graph to fp16.') |
||||
parser.add_argument('--input', required=True, help='Path to frozen graph.') |
||||
parser.add_argument('--output', required=True, help='Path to output graph.') |
||||
parser.add_argument('--ops', default=['Conv2D', 'MatMul'], nargs='+', |
||||
help='List of ops which weights are converted.') |
||||
args = parser.parse_args() |
||||
|
||||
DT_FLOAT = 1 |
||||
DT_HALF = 19 |
||||
|
||||
# For the frozen graphs, an every node that uses weights connected to Const nodes |
||||
# through an Identity node. Usually they're called in the same way with '/read' suffix. |
||||
# We'll replace all of them to Cast nodes. |
||||
|
||||
# Load the model |
||||
with tf.gfile.FastGFile(args.input) as f: |
||||
graph_def = tf.GraphDef() |
||||
graph_def.ParseFromString(f.read()) |
||||
|
||||
# Set of all inputs from desired nodes. |
||||
inputs = [] |
||||
for node in graph_def.node: |
||||
if node.op in args.ops: |
||||
inputs += node.input |
||||
|
||||
weightsNodes = [] |
||||
for node in graph_def.node: |
||||
# From the whole inputs we need to keep only an Identity nodes. |
||||
if node.name in inputs and node.op == 'Identity' and node.attr['T'].type == DT_FLOAT: |
||||
weightsNodes.append(node.input[0]) |
||||
|
||||
# Replace Identity to Cast. |
||||
node.op = 'Cast' |
||||
node.attr['DstT'].type = DT_FLOAT |
||||
node.attr['SrcT'].type = DT_HALF |
||||
del node.attr['T'] |
||||
del node.attr['_class'] |
||||
|
||||
# Convert weights to halfs. |
||||
for node in graph_def.node: |
||||
if node.name in weightsNodes: |
||||
node.attr['dtype'].type = DT_HALF |
||||
node.attr['value'].tensor.dtype = DT_HALF |
||||
|
||||
floats = node.attr['value'].tensor.tensor_content |
||||
|
||||
floats = struct.unpack('f' * (len(floats) / 4), floats) |
||||
halfs = np.array(floats).astype(np.float16).view(np.uint16) |
||||
node.attr['value'].tensor.tensor_content = struct.pack('H' * len(halfs), *halfs) |
||||
|
||||
tf.train.write_graph(graph_def, "", args.output, as_text=False) |
||||
Loading…
Reference in new issue