import time
from gluoncv import model_zoo, data, utils
import gluoncv as gcv
from gluoncv.utils import try_import_cv2
cv2 = try_import_cv2()
import mxnet as mx
#import cv2
# Load the model
net = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True)
#net = model_zoo.get_model('yolo3_darknet53_voc', pretrained=True)
# Compile the model for faster speed
net.hybridize()
#import cv2
def gstreamer_pipeline(
capture_width=1280,
capture_height=720,
display_width=1280,
display_height=720,
framerate=60,
flip_method=0,
):
return (
"nvarguscamerasrc ! "
"video/x-raw(memory:NVMM), "
"width=(int)%d, height=(int)%d, "
"format=(string)NV12, framerate=(fraction)%d/1 ! "
"nvvidconv flip-method=%d ! "
"video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! "
"videoconvert ! "
"video/x-raw, format=(string)BGR ! appsink"
% (
capture_width,
capture_height,
framerate,
flip_method,
display_width,
display_height,
)
)
cap = cv2.VideoCapture(gstreamer_pipeline(flip_method=0), cv2.CAP_GSTREAMER)
# Load the webcam handler
#cap = cv2.VideoCapture(0)
time.sleep(1) ### letting the camera autofocus
axes = None
NUM_FRAMES = 200 # you can change this
for i in range(NUM_FRAMES):
# Load frame from the camera
ret, frame = cap.read()
# Image pre-processing
frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')
rgb_nd, frame = gcv.data.transforms.presets.ssd.transform_test(frame, short=512, max_size=700)
# Run frame through network
class_IDs, scores, bounding_boxes = net(rgb_nd)
# Display the result
img = gcv.utils.viz.cv_plot_bbox(frame, bounding_boxes[0], scores[0], class_IDs[0], class_names=net.classes)
gcv.utils.viz.cv_plot_image(img)
# cv2.waitKey(1)
if cv2.waitKey(1) & 0xFF == ord('q'):
# 存储图片
cv2.imwrite("camera.jpeg", frame)
break
cap.release()
cv2.destroyAllWindows()
