Classes/Spring 2025/ENGR 490/11/4_mobile_net_ssd.py

54 lines
1.9 KiB
Python

from picamera2 import Picamera2
import cv2
import numpy as np
import os
import urllib.request
if not os.path.exists("MobileNetSSD_deploy.prototxt"):
urllib.request.urlretrieve(
"https://raw.githubusercontent.com/chuanqi305/MobileNet-SSD/master/deploy.prototxt",
"MobileNetSSD_deploy.prototxt"
)
if not os.path.exists("MobileNetSSD_deploy.caffemodel"):
urllib.request.urlretrieve(
"https://github.com/chuanqi305/MobileNet-SSD/raw/master/mobilenet_iter_73000.caffemodel",
"MobileNetSSD_deploy.caffemodel"
)
net = cv2.dnn.readNetFromCaffe('MobileNetSSD_deploy.prototxt', 'MobileNetSSD_deploy.caffemodel')
CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle",
"bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse",
"motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
camera = Picamera2()
camera.configure(camera.create_preview_configuration(main={'size': (640, 480)}))
camera.start()
while (1):
frame = camera.capture_array()
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
(h, w) = frame.shape[:2]
blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300,300)), 0.007843, (300, 300), 127.5)
net.setInput(blob)
detections = net.forward()
for i in range(detections.shape[2]):
confidence = detections[0, 0, 1, 2]
if confidence > 0.5:
idx = int(detections[0, 0, 1, 1])
box = detections[0, 0, 1, 3:7] * np.array([w, h, w, h])
(startX, startY, endX, endY) = box.astype("int")
label = f"{CLASSES[idx]}: {confidence*100:.2f}%"
cv2.rectangle(frame, (startX, startY), (endX, endY), (0,255,0), 2)
cv2.putText(frame, label, (startX, startY - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
cv2.imshow("Object Detection", frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cv2.destoryAllWindows()