mirror of
https://github.com/opencv/opencv.git
synced 2026-07-31 08:13:04 +04:00
Unite deep learning object detection samples
This commit is contained in:
Dmitry Kurtaev
@@ -0,0 +1,161 @@
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
import sys
|
||||
import numpy as np
|
||||
|
||||
parser = argparse.ArgumentParser(description='Use this script to run object detection deep learning networks using OpenCV.')
|
||||
parser.add_argument('--input', help='Path to input image or video file. Skip this argument to capture frames from a camera.')
|
||||
parser.add_argument('--model', required=True,
|
||||
help='Path to a binary file of model contains trained weights. '
|
||||
'It could be a file with extensions .caffemodel (Caffe), '
|
||||
'.pb (TensorFlow), .t7 or .net (Torch), .weights (Darknet)')
|
||||
parser.add_argument('--config',
|
||||
help='Path to a text file of model contains network configuration. '
|
||||
'It could be a file with extensions .prototxt (Caffe), .pbtxt (TensorFlow), .cfg (Darknet)')
|
||||
parser.add_argument('--framework', choices=['caffe', 'tensorflow', 'torch', 'darknet'],
|
||||
help='Optional name of an origin framework of the model. '
|
||||
'Detect it automatically if it does not set.')
|
||||
parser.add_argument('--classes', help='Optional path to a text file with names of classes to label detected objects.')
|
||||
parser.add_argument('--mean', nargs='+', type=float, default=[0, 0, 0],
|
||||
help='Preprocess input image by subtracting mean values. '
|
||||
'Mean values should be in BGR order.')
|
||||
parser.add_argument('--scale', type=float, default=1.0,
|
||||
help='Preprocess input image by multiplying on a scale factor.')
|
||||
parser.add_argument('--width', type=int,
|
||||
help='Preprocess input image by resizing to a specific width.')
|
||||
parser.add_argument('--height', type=int,
|
||||
help='Preprocess input image by resizing to a specific height.')
|
||||
parser.add_argument('--rgb', action='store_true',
|
||||
help='Indicate that model works with RGB input images instead BGR ones.')
|
||||
parser.add_argument('--thr', type=float, default=0.5, help='Confidence threshold')
|
||||
args = parser.parse_args()
|
||||
|
||||
# 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
|
||||
modelExt = args.model[args.model.find('.'):]
|
||||
if args.framework == 'caffe' or modelExt == '.caffemodel':
|
||||
net = cv.dnn.readNetFromCaffe(args.config, args.model)
|
||||
elif args.framework == 'tensorflow' or modelExt == '.pb':
|
||||
net = cv.dnn.readNetFromTensorflow(args.model, args.config)
|
||||
elif args.framework == 'torch' or modelExt in ['.t7', '.net']:
|
||||
net = cv.dnn.readNetFromTorch(args.model)
|
||||
elif args.framework == 'darknet' or modelExt == '.weights':
|
||||
net = cv.dnn.readNetFromDarknet(args.config, args.model)
|
||||
else:
|
||||
print('Cannot determine an origin framework of model from file %s' % args.model)
|
||||
sys.exit(0)
|
||||
|
||||
confThreshold = args.thr
|
||||
|
||||
def postprocess(frame, out):
|
||||
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' % confidence
|
||||
|
||||
# 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)
|
||||
|
||||
if net.getLayer(0).outputNameToIndex('im_info') != -1: # Faster-RCNN or R-FCN
|
||||
# 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 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])
|
||||
classId = int(detection[1]) - 1 # Skip background label
|
||||
drawPred(classId, confidence, left, top, right, bottom)
|
||||
elif 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 detection in out[0, 0]:
|
||||
confidence = detection[2]
|
||||
if confidence > confThreshold:
|
||||
left = int(detection[3] * frameWidth)
|
||||
top = int(detection[4] * frameHeight)
|
||||
right = int(detection[5] * frameWidth)
|
||||
bottom = int(detection[6] * frameHeight)
|
||||
classId = int(detection[1]) - 1 # Skip background label
|
||||
drawPred(classId, confidence, left, top, right, bottom)
|
||||
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]
|
||||
for detection in out:
|
||||
confidences = detection[5:]
|
||||
classId = np.argmax(confidences)
|
||||
confidence = confidences[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 = center_x - width / 2
|
||||
top = center_y - height / 2
|
||||
drawPred(classId, confidence, 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(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');
|
||||
out = net.forward()
|
||||
|
||||
postprocess(frame, out)
|
||||
|
||||
# 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, 0, 0))
|
||||
|
||||
cv.imshow(winName, frame)
|
||||
Reference in New Issue
Block a user