import cv2
from ultralytics import YOLO
from scipy.spatial import distance as dist
import datetime
import json

focal_length = 258

real_sizes = {
    'person': 1.65,
    'car': 1.5,
    'bicycle': 1.0,
    'dog': 0.5,
    'cat': 0.3,
    'truck': 2.0,
    'bus': 3.5,
    'motorcycle': 1.2,
    'sheep': 1.0,
    'horse': 1.5,
    'elephant': 3.0,
    'giraffe': 5.5,
    'table': 0.75,
    'chair': 0.45,
    'sofa': 0.85,
    'tv': 1.0,
    'refrigerator': 1.8,
    'door': 2.0,
    'window': 1.5,
    'tree': 5.0,
    'flower': 0.3,
    'basketball': 0.24,
    'soccer_ball': 0.22,
    'apple': 0.1,
    'banana': 0.2,
    'bottle': 0.25,
    'cup': 0.1,
    'pencil': 0.19,
    'book': 0.25,
    'laptop': 0.02,
    'backpack': 0.5,
    'stool': 0.75,
    'bench': 0.9,
    'kitchen_island': 0.9,
    'fire_hydrant': 1.0,
    'traffic_light': 2.5,
    'sign': 1.5,
    'cactus': 1.0,
    'palm_tree': 6.0,
    'fence': 1.5,
    'swing': 1.5,
    'ladder': 2.0,
    'scooter': 0.9,
    'skateboard': 0.1,
}

model = YOLO("yolo11n.pt")
cap = cv2.VideoCapture(0)

def get_label(class_id, yolo_classes):
    return yolo_classes[class_id]

while cap.isOpened():
    success, frame = cap.read()
    if success:
        results = model(frame)
        annotated_frame = results[0].plot()
        detection_data = []
        for result in results:
            for box in result.boxes:
                x1, y1, x2, y2 = box.xyxy[0]
                confidence = box.conf[0]
                label = model.names[int(box.cls[0])]  
                class_id = result.boxes.cls[0]
                name = get_label(int(class_id), model.names)
                box_size = dist.euclidean((x1, y1), (x2, y2))

                real_size = real_sizes.get(label, 1.0)
                distance = (real_size * focal_length) / box_size

                detection_data.append({
                    "time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "class_id": int(class_id),
                    'label': label,
                    "name": name,
                    'confidence': confidence.item(),
                    'box': [int(x1), int(y1), int(x2), int(y2)],
                    "x1": int(x1),
                    "y1": int(y1),
                    "x2": int(x2),
                    "y2": int(y2),
                    "distance": distance
                })

        print(json.dumps(detection_data))

        for detection in detection_data:
            x1, y1, x2, y2 = detection['box']
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            text = f"{detection['label']}: {detection['confidence']:.2f}, Distance: {detection['distance']:.2f}"
            cv2.putText(frame, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            
        cv2.imshow("Prediksi Kamera", frame)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        break

cap.release()
cv2.destroyAllWindows()