akshaikrishna's picture
Inital commit
9c0a13e
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()