mirror of https://github.com/opencv/opencv.git
Open Source Computer Vision Library
https://opencv.org/
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
204 lines
9.0 KiB
204 lines
9.0 KiB
import cv2 as cv |
|
import argparse |
|
import numpy as np |
|
|
|
from common import * |
|
from tf_text_graph_common import readTextMessage |
|
from tf_text_graph_ssd import createSSDGraph |
|
from tf_text_graph_faster_rcnn import createFasterRCNNGraph |
|
|
|
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_HALIDE, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE, cv.dnn.DNN_BACKEND_OPENCV) |
|
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD) |
|
|
|
parser = argparse.ArgumentParser(add_help=False) |
|
parser.add_argument('--zoo', default=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models.yml'), |
|
help='An optional path to file with preprocessing parameters.') |
|
parser.add_argument('--input', help='Path to input image or video file. Skip this argument to capture frames from a camera.') |
|
parser.add_argument('--out_tf_graph', default='graph.pbtxt', |
|
help='For models from TensorFlow Object Detection API, you may ' |
|
'pass a .config file which was used for training through --config ' |
|
'argument. This way an additional .pbtxt file with TensorFlow graph will be created.') |
|
parser.add_argument('--framework', choices=['caffe', 'tensorflow', 'torch', 'darknet', 'dldt'], |
|
help='Optional name of an origin framework of the model. ' |
|
'Detect it automatically if it does not set.') |
|
parser.add_argument('--thr', type=float, default=0.5, help='Confidence threshold') |
|
parser.add_argument('--nms', type=float, default=0.4, help='Non-maximum suppression threshold') |
|
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int, |
|
help="Choose one of computation backends: " |
|
"%d: automatically (by default), " |
|
"%d: Halide language (http://halide-lang.org/), " |
|
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), " |
|
"%d: OpenCV implementation" % backends) |
|
parser.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int, |
|
help='Choose one of target computation devices: ' |
|
'%d: CPU target (by default), ' |
|
'%d: OpenCL, ' |
|
'%d: OpenCL fp16 (half-float precision), ' |
|
'%d: VPU' % targets) |
|
args, _ = parser.parse_known_args() |
|
add_preproc_args(args.zoo, parser, 'object_detection') |
|
parser = argparse.ArgumentParser(parents=[parser], |
|
description='Use this script to run object detection deep learning networks using OpenCV.', |
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
|
args = parser.parse_args() |
|
|
|
args.model = findFile(args.model) |
|
args.config = findFile(args.config) |
|
args.classes = findFile(args.classes) |
|
|
|
# If config specified, try to load it as TensorFlow Object Detection API's pipeline. |
|
config = readTextMessage(args.config) |
|
if 'model' in config: |
|
print('TensorFlow Object Detection API config detected') |
|
if 'ssd' in config['model'][0]: |
|
print('Preparing text graph representation for SSD model: ' + args.out_tf_graph) |
|
createSSDGraph(args.model, args.config, args.out_tf_graph) |
|
args.config = args.out_tf_graph |
|
elif 'faster_rcnn' in config['model'][0]: |
|
print('Preparing text graph representation for Faster-RCNN model: ' + args.out_tf_graph) |
|
createFasterRCNNGraph(args.model, args.config, args.out_tf_graph) |
|
args.config = args.out_tf_graph |
|
|
|
|
|
# Load names of classes |
|
classes = None |
|
if args.classes: |
|
with open(args.classes, 'rt') as f: |
|
classes = f.read().rstrip('\n').split('\n') |
|
|
|
# Load a network |
|
net = cv.dnn.readNet(cv.samples.findFile(args.model), cv.samples.findFile(args.config), args.framework) |
|
net.setPreferableBackend(args.backend) |
|
net.setPreferableTarget(args.target) |
|
outNames = net.getUnconnectedOutLayersNames() |
|
|
|
confThreshold = args.thr |
|
nmsThreshold = args.nms |
|
|
|
def postprocess(frame, outs): |
|
frameHeight = frame.shape[0] |
|
frameWidth = frame.shape[1] |
|
|
|
def drawPred(classId, conf, left, top, right, bottom): |
|
# Draw a bounding box. |
|
cv.rectangle(frame, (left, top), (right, bottom), (0, 255, 0)) |
|
|
|
label = '%.2f' % conf |
|
|
|
# Print a label of class. |
|
if classes: |
|
assert(classId < len(classes)) |
|
label = '%s: %s' % (classes[classId], label) |
|
|
|
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1) |
|
top = max(top, labelSize[1]) |
|
cv.rectangle(frame, (left, top - labelSize[1]), (left + labelSize[0], top + baseLine), (255, 255, 255), cv.FILLED) |
|
cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0)) |
|
|
|
layerNames = net.getLayerNames() |
|
lastLayerId = net.getLayerId(layerNames[-1]) |
|
lastLayer = net.getLayer(lastLayerId) |
|
|
|
classIds = [] |
|
confidences = [] |
|
boxes = [] |
|
if lastLayer.type == 'DetectionOutput': |
|
# Network produces output blob with a shape 1x1xNx7 where N is a number of |
|
# detections and an every detection is a vector of values |
|
# [batchId, classId, confidence, left, top, right, bottom] |
|
for out in outs: |
|
for detection in out[0, 0]: |
|
confidence = detection[2] |
|
if confidence > confThreshold: |
|
left = int(detection[3]) |
|
top = int(detection[4]) |
|
right = int(detection[5]) |
|
bottom = int(detection[6]) |
|
width = right - left + 1 |
|
height = bottom - top + 1 |
|
if width * height <= 1: |
|
left = int(detection[3] * frameWidth) |
|
top = int(detection[4] * frameHeight) |
|
right = int(detection[5] * frameWidth) |
|
bottom = int(detection[6] * frameHeight) |
|
width = right - left + 1 |
|
height = bottom - top + 1 |
|
classIds.append(int(detection[1]) - 1) # Skip background label |
|
confidences.append(float(confidence)) |
|
boxes.append([left, top, width, height]) |
|
elif lastLayer.type == 'Region': |
|
# Network produces output blob with a shape NxC where N is a number of |
|
# detected objects and C is a number of classes + 4 where the first 4 |
|
# numbers are [center_x, center_y, width, height] |
|
classIds = [] |
|
confidences = [] |
|
boxes = [] |
|
for out in outs: |
|
for detection in out: |
|
scores = detection[5:] |
|
classId = np.argmax(scores) |
|
confidence = scores[classId] |
|
if confidence > confThreshold: |
|
center_x = int(detection[0] * frameWidth) |
|
center_y = int(detection[1] * frameHeight) |
|
width = int(detection[2] * frameWidth) |
|
height = int(detection[3] * frameHeight) |
|
left = int(center_x - width / 2) |
|
top = int(center_y - height / 2) |
|
classIds.append(classId) |
|
confidences.append(float(confidence)) |
|
boxes.append([left, top, width, height]) |
|
else: |
|
print('Unknown output layer type: ' + lastLayer.type) |
|
exit() |
|
|
|
indices = cv.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold) |
|
for i in indices: |
|
i = i[0] |
|
box = boxes[i] |
|
left = box[0] |
|
top = box[1] |
|
width = box[2] |
|
height = box[3] |
|
drawPred(classIds[i], confidences[i], left, top, left + width, top + height) |
|
|
|
# Process inputs |
|
winName = 'Deep learning object detection in OpenCV' |
|
cv.namedWindow(winName, cv.WINDOW_NORMAL) |
|
|
|
def callback(pos): |
|
global confThreshold |
|
confThreshold = pos / 100.0 |
|
|
|
cv.createTrackbar('Confidence threshold, %', winName, int(confThreshold * 100), 99, callback) |
|
|
|
cap = cv.VideoCapture(cv.samples.findFileOrKeep(args.input) if args.input else 0) |
|
while cv.waitKey(1) < 0: |
|
hasFrame, frame = cap.read() |
|
if not hasFrame: |
|
cv.waitKey() |
|
break |
|
|
|
frameHeight = frame.shape[0] |
|
frameWidth = frame.shape[1] |
|
|
|
# Create a 4D blob from a frame. |
|
inpWidth = args.width if args.width else frameWidth |
|
inpHeight = args.height if args.height else frameHeight |
|
blob = cv.dnn.blobFromImage(frame, args.scale, (inpWidth, inpHeight), args.mean, args.rgb, crop=False) |
|
|
|
# Run a model |
|
net.setInput(blob) |
|
if net.getLayer(0).outputNameToIndex('im_info') != -1: # Faster-RCNN or R-FCN |
|
frame = cv.resize(frame, (inpWidth, inpHeight)) |
|
net.setInput(np.array([[inpHeight, inpWidth, 1.6]], dtype=np.float32), 'im_info') |
|
outs = net.forward(outNames) |
|
|
|
postprocess(frame, outs) |
|
|
|
# Put efficiency information. |
|
t, _ = net.getPerfProfile() |
|
label = 'Inference time: %.2f ms' % (t * 1000.0 / cv.getTickFrequency()) |
|
cv.putText(frame, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0)) |
|
|
|
cv.imshow(winName, frame)
|
|
|