|
import cv2 |
|
import numpy as np |
|
import gradio as gr |
|
|
|
net = cv2.dnn.readNetFromCaffe( |
|
"MobileNetSSD_deploy.prototxt", "MobileNetSSD_deploy.caffemodel" |
|
) |
|
|
|
class_names = [ |
|
"background", |
|
"aeroplane", |
|
"bicycle", |
|
"bird", |
|
"boat", |
|
"bottle", |
|
"bus", |
|
"car", |
|
"cat", |
|
"chair", |
|
"cow", |
|
"diningtable", |
|
"dog", |
|
"horse", |
|
"motorbike", |
|
"person", |
|
"pottedplant", |
|
"sheep", |
|
"sofa", |
|
"train", |
|
"tvmonitor", |
|
] |
|
|
|
|
|
def detect_objects(image): |
|
frame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) |
|
blob = cv2.dnn.blobFromImage( |
|
frame, 0.007843, (300, 300), (127.5, 127.5, 127.5), swapRB=False, crop=False |
|
) |
|
net.setInput(blob) |
|
detections = net.forward() |
|
|
|
for i in range(detections.shape[2]): |
|
confidence = detections[0, 0, i, 2] |
|
if confidence > 0.2: |
|
idx = int(detections[0, 0, i, 1]) |
|
box = detections[0, 0, i, 3:7] * np.array( |
|
[frame.shape[1], frame.shape[0], frame.shape[1], frame.shape[0]] |
|
) |
|
(startX, startY, endX, endY) = box.astype("int") |
|
|
|
label = f"{class_names[idx]}: {confidence:.2f}" |
|
cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2) |
|
y = startY - 15 if startY - 15 > 15 else startY + 15 |
|
cv2.putText( |
|
frame, |
|
label, |
|
(startX, y), |
|
cv2.FONT_HERSHEY_SIMPLEX, |
|
0.5, |
|
(255, 255, 255), |
|
2, |
|
) |
|
|
|
return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
|
|
|
|
|
iface = gr.Interface(fn=detect_objects, inputs="image", outputs="image", live=True) |
|
iface.launch() |
|
|