86 lines
2.1 KiB
Python
86 lines
2.1 KiB
Python
# backend\ml\predict_detector_model.py
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
from ultralytics import YOLO
|
|
|
|
|
|
def clamp01(v):
|
|
return max(0.0, min(1.0, float(v)))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
parser.add_argument("--image", required=True)
|
|
parser.add_argument("--conf", type=float, default=0.30)
|
|
parser.add_argument("--imgsz", type=int, default=640)
|
|
args = parser.parse_args()
|
|
|
|
root = Path(args.root)
|
|
image_path = Path(args.image)
|
|
|
|
model_path = root / "detector" / "model" / "best.pt"
|
|
if not model_path.exists():
|
|
print(json.dumps({
|
|
"available": False,
|
|
"source": "detector_missing",
|
|
"boxes": [],
|
|
}))
|
|
return
|
|
|
|
img = Image.open(image_path).convert("RGB")
|
|
img_w, img_h = img.size
|
|
|
|
model = YOLO(str(model_path))
|
|
results = model.predict(
|
|
source=str(image_path),
|
|
conf=args.conf,
|
|
imgsz=args.imgsz,
|
|
verbose=False,
|
|
device=0 if __import__("torch").cuda.is_available() else "cpu",
|
|
)
|
|
|
|
boxes = []
|
|
|
|
if results:
|
|
r = results[0]
|
|
names = r.names or {}
|
|
|
|
if r.boxes is not None:
|
|
for b in r.boxes:
|
|
cls_id = int(b.cls[0].item())
|
|
score = float(b.conf[0].item())
|
|
label = str(names.get(cls_id, cls_id))
|
|
|
|
x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]
|
|
|
|
x = clamp01(x1 / img_w)
|
|
y = clamp01(y1 / img_h)
|
|
w = clamp01((x2 - x1) / img_w)
|
|
h = clamp01((y2 - y1) / img_h)
|
|
|
|
if w <= 0 or h <= 0:
|
|
continue
|
|
|
|
boxes.append({
|
|
"label": label,
|
|
"score": score,
|
|
"x": x,
|
|
"y": y,
|
|
"w": w,
|
|
"h": h,
|
|
})
|
|
|
|
print(json.dumps({
|
|
"available": True,
|
|
"source": "yolo_detector",
|
|
"boxes": boxes,
|
|
}, ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |