File size: 3,211 Bytes
92f31cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | 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.")
|