Update MobileNet object detection sample

pull/9333/head
Dmitry Kurtaev 8 years ago
parent b67c64e2c8
commit 1e3052d3cd
  1. 31
      samples/dnn/mobilenet_ssd_python.py
  2. 6
      samples/dnn/ssd_mobilenet_object_detection.cpp

@ -23,24 +23,25 @@ classNames = ('background',
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--video", help="path to video file. If empty, camera's stream will be used")
parser.add_argument("--prototxt", default="MobileNetSSD_300x300.prototxt",
parser.add_argument("--prototxt", default="MobileNetSSD_deploy.prototxt",
help="path to caffe prototxt")
parser.add_argument("-c", "--caffemodel", help="path to caffemodel file, download it here: "
"https://github.com/chuanqi305/MobileNet-SSD/blob/master/MobileNetSSD_train.caffemodel")
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")
args = parser.parse_args()
net = dnn.readNetFromCaffe(args.prototxt, args.caffemodel)
net = cv.dnn.readNetFromCaffe(args.prototxt, args.caffemodel)
if len(args.video):
cap = cv2.VideoCapture(args.video)
cap = cv.VideoCapture(args.video)
else:
cap = cv2.VideoCapture(0)
cap = cv.VideoCapture(0)
while True:
# Capture frame-by-frame
ret, frame = cap.read()
blob = dnn.blobFromImage(frame, inScaleFactor, (inWidth, inHeight), meanVal)
blob = cv.dnn.blobFromImage(frame, inScaleFactor, (inWidth, inHeight), meanVal)
net.setInput(blob)
detections = net.forward()
@ -71,17 +72,17 @@ if __name__ == "__main__":
xRightTop = int(detections[0, 0, i, 5] * cols)
yRightTop = int(detections[0, 0, i, 6] * rows)
cv2.rectangle(frame, (xLeftBottom, yLeftBottom), (xRightTop, yRightTop),
cv.rectangle(frame, (xLeftBottom, yLeftBottom), (xRightTop, yRightTop),
(0, 255, 0))
label = classNames[class_id] + ": " + str(confidence)
labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
cv2.rectangle(frame, (xLeftBottom, yLeftBottom - labelSize[1]),
cv.rectangle(frame, (xLeftBottom, yLeftBottom - labelSize[1]),
(xLeftBottom + labelSize[0], yLeftBottom + baseLine),
(255, 255, 255), cv2.FILLED)
cv2.putText(frame, label, (xLeftBottom, yLeftBottom),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
(255, 255, 255), cv.FILLED)
cv.putText(frame, label, (xLeftBottom, yLeftBottom),
cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0))
cv2.imshow("detections", frame)
if cv2.waitKey(1) >= 0:
cv.imshow("detections", frame)
if cv.waitKey(1) >= 0:
break

@ -27,12 +27,12 @@ const char* about = "This sample uses Single-Shot Detector "
"(https://arxiv.org/abs/1512.02325)"
"to detect objects on image.\n"
".caffemodel model's file is avaliable here: "
"https://github.com/chuanqi305/MobileNet-SSD/blob/master/MobileNetSSD_train.caffemodel\n";
"https://github.com/chuanqi305/MobileNet-SSD\n";
const char* params
= "{ help | false | print usage }"
"{ proto | MobileNetSSD_300x300.prototxt | model configuration }"
"{ model | | model weights }"
"{ proto | MobileNetSSD_deploy.prototxt | model configuration }"
"{ model | MobileNetSSD_deploy.caffemodel | model weights }"
"{ video | | video for detection }"
"{ out | | path to output video file}"
"{ min_confidence | 0.2 | min confidence }";

Loading…
Cancel
Save