Spaces:
Sleeping
Sleeping
File size: 4,976 Bytes
70ebc90 | 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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | from ultralytics import YOLO
import cv2
from pathlib import Path
import numpy as np
import os
# Get project root
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# ---------------- USER CONFIG ----------------
MODEL_PATH = os.path.join(PROJECT_ROOT, "models/model_a.pt") # YOLOv8 model path
DEVICE = "cpu" # or "0" for GPU
IMG_SIZE = 640
CONF_THR = 0.25
IOU_THR = 0.5
OUTPUT_DIR = os.path.join(PROJECT_ROOT, "outputs/3_detection") # annotated images folder
LABEL_DIR = os.path.join(PROJECT_ROOT, "outputs/3_detection/labels") # save labels folder
# --------------------------------------------
def ensure_dir(p):
Path(p).mkdir(parents=True, exist_ok=True)
def xyxy_to_xywh(xyxy):
"""Convert (x1,y1,x2,y2) to (cx,cy,w,h)"""
x1, y1, x2, y2 = xyxy
w = x2 - x1
h = y2 - y1
cx = x1 + w/2
cy = y1 + h/2
return cx, cy, w, h
def non_max_suppression(boxes, scores, iou_thresh=0.5):
"""Simple NMS to remove overlapping boxes"""
if len(boxes) == 0:
return []
idxs = cv2.dnn.NMSBoxes(
bboxes=[[int(x1), int(y1), int(x2-x1), int(y2-y1)] for (x1,y1,x2,y2) in boxes],
scores=[float(s) for s in scores],
score_threshold=0.0,
nms_threshold=iou_thresh
)
if len(idxs) == 0:
return []
return idxs.flatten().tolist()
def save_labels_yolo_format(image_path, hbb_list, label_dir=LABEL_DIR):
"""Saves labels in YOLO TXT format and returns path"""
ensure_dir(label_dir)
stem = Path(image_path).stem
label_path = Path(label_dir) / f"{stem}.txt"
img = cv2.imread(str(image_path))
h, w = img.shape[:2]
with open(label_path, "w") as f:
for det in hbb_list:
cls_id = det["class_id"]
cx, cy, bw, bh = det["xywh"]
cxn, cyn, w_n, h_n = cx/w, cy/h, bw/w, bh/h
f.write(f"{cls_id} {cxn:.6f} {cyn:.6f} {w_n:.6f} {h_n:.6f}\n")
return str(label_path)
def run_modelA(image_path, model_path=MODEL_PATH, device=DEVICE, imgsz=IMG_SIZE,
conf_thr=CONF_THR, save_annotated=True, outdir=OUTPUT_DIR,
iou_thr=IOU_THR, save_labels=True, model=None):
"""
Runs YOLOv8 HBB detection on image and optionally saves labels.
Args:
model (ultralytics.YOLO, optional): Preloaded YOLO model instance.
Returns:
hbb_list (list of dicts): Each detection info
label_path (str or None): Path to saved label file
"""
ensure_dir(outdir)
ensure_dir(LABEL_DIR)
img_bgr = cv2.imread(str(image_path))
if img_bgr is None:
print(f"❌ Could not read image: {image_path}")
return [], None
# ⚠️ No resizing here, image is already 640×640 from capture.py
model = model if model is not None else YOLO(model_path)
results = model.predict(source=img_bgr, device=device, imgsz=imgsz, conf=conf_thr, verbose=False)
result = results[0]
try:
names = model.names
except Exception:
names = {i: str(i) for i in range(1000)}
hbb_list = []
label_path = None
if len(result.boxes) > 0:
xyxy_all = result.boxes.xyxy.cpu().numpy()
conf_all = result.boxes.conf.cpu().numpy()
cls_all = result.boxes.cls.cpu().numpy().astype(int)
keep_idxs = non_max_suppression(xyxy_all, conf_all, iou_thresh=iou_thr)
for i in keep_idxs:
x1, y1, x2, y2 = xyxy_all[i]
cx, cy, w_box, h_box = xyxy_to_xywh((x1, y1, x2, y2))
conf = float(conf_all[i])
cls_id = int(cls_all[i])
cls_name = names.get(cls_id, str(cls_id))
entry = {
"class_id": cls_id,
"class_name": cls_name,
"conf": conf,
"xyxy": (float(x1), float(y1), float(x2), float(y2)),
"xywh": (float(cx), float(cy), float(w_box), float(h_box))
}
hbb_list.append(entry)
# Annotate image
if save_annotated and len(hbb_list) > 0:
annotated = result.plot()
outpath = Path(outdir) / (Path(image_path).stem + "_annotated.jpg")
cv2.imwrite(str(outpath), cv2.cvtColor(annotated, cv2.COLOR_RGB2BGR))
# Save labels
if save_labels and len(hbb_list) > 0:
label_path = save_labels_yolo_format(image_path, hbb_list, label_dir=LABEL_DIR)
return hbb_list, label_path
# ----------------- Test run -----------------
if __name__ == "__main__":
IMAGE_PATH = "outputs/1_captured_images/capture_1757686922_0.jpg" # replace with your test image
detections, label_path = run_modelA(IMAGE_PATH)
if detections:
print("\nDetected objects:")
for i, det in enumerate(detections):
print(f"{i+1}. {det['class_name']} conf={det['conf']:.3f}, xyxy={tuple(round(v,2) for v in det['xyxy'])}")
print(f"\nLabels saved at: {label_path}")
else:
print("No objects detected.")
|