import cv2
import numpy as np
from ultralytics import YOLO

# ---------- CONFIG ----------
MODEL_PATH = r"R:\computer_vison\Trained_models\best.pt"
CONF_THRESH = 0.55  
ASSUMED_HUMAN_HEIGHT_M = 1.70
WINDOW_NAME = "Stone Virtual Measurements"
FRAME_SCALE = 1.0
SHARPEN = True
# -----------------------------


model = YOLO(MODEL_PATH)
print(f"Model loaded from: {MODEL_PATH}")


names = model.names.copy()
print(f"Original model class names: {names}")


fixed_names = {0: "human", 1: "stone"}
print(f"Using fixed label mapping: {fixed_names}")


cap = cv2.VideoCapture(0)
if not cap.isOpened():
    raise RuntimeError("No camera found! Check webcam connection.")

cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
cv2.resizeWindow(WINDOW_NAME, 960, 720)

print("Starting real-time detection... Press ESC to exit.")

while True:
    ret, frame = cap.read()
    if not ret:
        print("Camera frame not received.")
        break

    
    if FRAME_SCALE != 1.0:
        frame = cv2.resize(frame, None, fx=FRAME_SCALE, fy=FRAME_SCALE, interpolation=cv2.INTER_LINEAR)

   
    if SHARPEN:
        kernel = np.array([[0, -1, 0],
                           [-1, 5, -1],
                           [0, -1, 0]], np.float32)
        frame = cv2.filter2D(frame, -1, kernel)

   
    results = model.predict(source=frame, conf=CONF_THRESH, verbose=False)
    boxes = results[0].boxes.xyxy.cpu().numpy()
    classes = results[0].boxes.cls.cpu().numpy().astype(int)

    human_height_px = None
    stone_box = None

    for i, box in enumerate(boxes):
        x1, y1, x2, y2 = box.astype(int)
        cls = classes[i]

        label = fixed_names.get(cls, f"class{cls}")
        color = (0, 255, 0) if label == "human" else (0, 255, 255)

        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
        cv2.putText(frame, label, (x1, max(30, y1 - 10)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

        if label == "human":
            human_height_px = y2 - y1
        elif label == "stone":
            stone_box = (x1, y1, x2, y2)

    
    if human_height_px and stone_box:
        px_per_meter = human_height_px / ASSUMED_HUMAN_HEIGHT_M
        x1, y1, x2, y2 = stone_box

        stone_w_px = x2 - x1
        stone_h_px = y2 - y1
        stone_w_m = stone_w_px / px_per_meter
        stone_h_m = stone_h_px / px_per_meter
        stone_area_mm2 = (stone_w_m * stone_h_m) * 1e6

        cv2.putText(frame, f"Stone: {stone_w_m*1000:.1f} x {stone_h_m*1000:.1f} mm",
                    (x1, y2 + 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
        cv2.putText(frame, f"Area: {stone_area_mm2:.0f} mm^2",
                    (x1, y2 + 50), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)

    elif stone_box:
        cv2.putText(frame, "No human reference - scale unknown",
                    (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

    cv2.imshow(WINDOW_NAME, frame)

    
    cv2.resizeWindow(WINDOW_NAME, 960, 720)

    key = cv2.waitKey(1) & 0xFF
    if key == 27:  # ESC key
        break

cap.release()
cv2.destroyAllWindows()
print("Exited cleanly.")