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()