potok/opencv
1
0
Fork 0
mirror of https://github.com/opencv/opencv.git synced 2026-07-30 15:53:03 +04:00
Code Packages Releases Wiki Activity

Text TensorFlow graphs parsing. MobileNet-SSD for 90 classes.

Browse Source
This commit is contained in:
Dmitry Kurtaev
2017-09-28 16:51:47 +03:00
parent 8ac2c5d620
commit e4aa39f9e5
10 changed files with 536 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/dnn/mobilenet_ssd_accuracy.py
+131
View File
@@ -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()
samples/dnn/mobilenet_ssd_python.py
+67 -23
View File
@@ -1,3 +1,14 @@
# This script is used to demonstrate MobileNet-SSD network using OpenCV deep learning module.
#
# It works with model taken from https://github.com/chuanqi305/MobileNet-SSD/ that
# was trained in Caffe-SSD framework, https://github.com/weiliu89/caffe/tree/ssd.
# Model detects objects from 20 classes.
#
# Also TensorFlow model from TensorFlow object detection model zoo may be used to
# detect objects from 90 classes:
# https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md
# Text graph definition must be taken from opencv_extra:
# https://github.com/opencv/opencv_extra/tree/master/testdata/dnn/ssd_mobilenet_v1_coco.pbtxt
import numpy as np
import argparse
@@ -13,27 +24,58 @@ WHRatio = inWidth / float(inHeight)
inScaleFactor = 0.007843
meanVal = 127.5
classNames = ('background',
'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair',
'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant',
'sheep', 'sofa', 'train', 'tvmonitor')
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser = argparse.ArgumentParser(
description='Script to run MobileNet-SSD object detection network '
'trained either in Caffe or TensorFlow frameworks.')
parser.add_argument("--video", help="path to video file. If empty, camera's stream will be used")
parser.add_argument("--prototxt", default="MobileNetSSD_deploy.prototxt",
help="path to caffe prototxt")
parser.add_argument("-c", "--caffemodel", default="MobileNetSSD_deploy.caffemodel",
help="path to caffemodel file, download it here: "
"https://github.com/chuanqi305/MobileNet-SSD/")
parser.add_argument("--thr", default=0.2, help="confidence threshold to filter out weak detections")
help='Path to text network file: '
'MobileNetSSD_deploy.prototxt for Caffe model or '
'ssd_mobilenet_v1_coco.pbtxt from opencv_extra for TensorFlow model')
parser.add_argument("--weights", default="MobileNetSSD_deploy.caffemodel",
help='Path to weights: '
'MobileNetSSD_deploy.caffemodel for Caffe model or '
'frozen_inference_graph.pb from TensorFlow.')
parser.add_argument("--num_classes", default=20, type=int,
help="Number of classes. It's 20 for Caffe model from "
"https://github.com/chuanqi305/MobileNet-SSD/ and 90 for "
"TensorFlow model from https://github.com/tensorflow/models/tree/master/research/object_detection")
parser.add_argument("--thr", default=0.2, type=float, help="confidence threshold to filter out weak detections")
args = parser.parse_args()
net = cv.dnn.readNetFromCaffe(args.prototxt, args.caffemodel)
if args.num_classes == 20:
net = cv.dnn.readNetFromCaffe(args.prototxt, args.weights)
swapRB = False
classNames = { 0: 'background',
1: 'aeroplane', 2: 'bicycle', 3: 'bird', 4: 'boat',
5: 'bottle', 6: 'bus', 7: 'car', 8: 'cat', 9: 'chair',
10: 'cow', 11: 'diningtable', 12: 'dog', 13: 'horse',
14: 'motorbike', 15: 'person', 16: 'pottedplant',
17: 'sheep', 18: 'sofa', 19: 'train', 20: 'tvmonitor' }
else:
assert(args.num_classes == 90)
net = cv.dnn.readNetFromTensorflow(args.weights, args.prototxt)
swapRB = True
classNames = { 0: 'background',
1: 'person', 2: 'bicycle', 3: 'car', 4: 'motorcycle', 5: 'airplane', 6: 'bus',
7: 'train', 8: 'truck', 9: 'boat', 10: 'traffic light', 11: 'fire hydrant',
13: 'stop sign', 14: 'parking meter', 15: 'bench', 16: 'bird', 17: 'cat',
18: 'dog', 19: 'horse', 20: 'sheep', 21: 'cow', 22: 'elephant', 23: 'bear',
24: 'zebra', 25: 'giraffe', 27: 'backpack', 28: 'umbrella', 31: 'handbag',
32: 'tie', 33: 'suitcase', 34: 'frisbee', 35: 'skis', 36: 'snowboard',
37: 'sports ball', 38: 'kite', 39: 'baseball bat', 40: 'baseball glove',
41: 'skateboard', 42: 'surfboard', 43: 'tennis racket', 44: 'bottle',
46: 'wine glass', 47: 'cup', 48: 'fork', 49: 'knife', 50: 'spoon',
51: 'bowl', 52: 'banana', 53: 'apple', 54: 'sandwich', 55: 'orange',
56: 'broccoli', 57: 'carrot', 58: 'hot dog', 59: 'pizza', 60: 'donut',
61: 'cake', 62: 'chair', 63: 'couch', 64: 'potted plant', 65: 'bed',
67: 'dining table', 70: 'toilet', 72: 'tv', 73: 'laptop', 74: 'mouse',
75: 'remote', 76: 'keyboard', 77: 'cell phone', 78: 'microwave', 79: 'oven',
80: 'toaster', 81: 'sink', 82: 'refrigerator', 84: 'book', 85: 'clock',
86: 'vase', 87: 'scissors', 88: 'teddy bear', 89: 'hair drier', 90: 'toothbrush' }
if len(args.video):
if args.video:
cap = cv.VideoCapture(args.video)
else:
cap = cv.VideoCapture(0)
@@ -41,7 +83,7 @@ if __name__ == "__main__":
while True:
# Capture frame-by-frame
ret, frame = cap.read()
blob = cv.dnn.blobFromImage(frame, inScaleFactor, (inWidth, inHeight), meanVal, False)
blob = cv.dnn.blobFromImage(frame, inScaleFactor, (inWidth, inHeight), (meanVal, meanVal, meanVal), swapRB)
net.setInput(blob)
detections = net.forward()
@@ -74,14 +116,16 @@ if __name__ == "__main__":
cv.rectangle(frame, (xLeftBottom, yLeftBottom), (xRightTop, yRightTop),
(0, 255, 0))
label = classNames[class_id] + ": " + str(confidence)
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
if class_id in classNames:
label = classNames[class_id] + ": " + str(confidence)
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
cv.rectangle(frame, (xLeftBottom, yLeftBottom - labelSize[1]),
(xLeftBottom + labelSize[0], yLeftBottom + baseLine),
(255, 255, 255), cv.FILLED)
cv.putText(frame, label, (xLeftBottom, yLeftBottom),
cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
yLeftBottom = max(yLeftBottom, labelSize[1])
cv.rectangle(frame, (xLeftBottom, yLeftBottom - labelSize[1]),
(xLeftBottom + labelSize[0], yLeftBottom + baseLine),
(255, 255, 255), cv.FILLED)
cv.putText(frame, label, (xLeftBottom, yLeftBottom),
cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
cv.imshow("detections", frame)
if cv.waitKey(1) >= 0:
samples/dnn/shrink_tf_graph_weights.py
+62
View File
@@ -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)