| | import cv2
|
| | import numpy as np
|
| | from ultralytics import YOLO
|
| |
|
| |
|
| | 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:
|
| | break
|
| |
|
| | cap.release()
|
| | cv2.destroyAllWindows()
|
| | print("Exited cleanly.")
|
| | |
