123 lines
3.3 KiB
Python
123 lines
3.3 KiB
Python
# backend\ml\predict_detector_model.py
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
from ultralytics import YOLO
|
|
import torch
|
|
|
|
|
|
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)
|
|
parser.add_argument("--debug-image", action="store_true")
|
|
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",
|
|
"modelPath": str(model_path),
|
|
"boxes": [],
|
|
}, ensure_ascii=False))
|
|
return
|
|
|
|
img = Image.open(image_path).convert("RGB")
|
|
img_w, img_h = img.size
|
|
|
|
try:
|
|
model = YOLO(str(model_path))
|
|
except Exception as e:
|
|
print(json.dumps({
|
|
"available": False,
|
|
"source": "detector_load_failed",
|
|
"modelPath": str(model_path),
|
|
"error": repr(e),
|
|
"boxes": [],
|
|
}, ensure_ascii=False))
|
|
return
|
|
|
|
device = 0 if torch.cuda.is_available() else "cpu"
|
|
|
|
results = model.predict(
|
|
source=str(image_path),
|
|
conf=float(args.conf),
|
|
imgsz=int(args.imgsz),
|
|
verbose=False,
|
|
device=device,
|
|
)
|
|
|
|
boxes = []
|
|
model_names = {}
|
|
|
|
raw_box_count = 0
|
|
|
|
if results:
|
|
r = results[0]
|
|
model_names = {str(k): str(v) for k, v in (r.names or {}).items()}
|
|
|
|
if r.boxes is not None:
|
|
raw_box_count = len(r.boxes)
|
|
|
|
for b in r.boxes:
|
|
cls_id = int(b.cls[0].item())
|
|
score = float(b.conf[0].item())
|
|
label = str((r.names or {}).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,
|
|
})
|
|
|
|
if args.debug_image:
|
|
debug_dir = root / "detector" / "debug"
|
|
debug_dir.mkdir(parents=True, exist_ok=True)
|
|
out_path = debug_dir / f"{image_path.stem}_conf_{args.conf}.jpg"
|
|
plotted = r.plot()
|
|
Image.fromarray(plotted).save(out_path)
|
|
|
|
print(json.dumps({
|
|
"available": True,
|
|
"source": "yolo_detector",
|
|
"modelPath": str(model_path),
|
|
"image": str(image_path),
|
|
"conf": float(args.conf),
|
|
"imgsz": int(args.imgsz),
|
|
"device": str(device),
|
|
"imageWidth": img_w,
|
|
"imageHeight": img_h,
|
|
"classNames": model_names,
|
|
"rawBoxCount": raw_box_count,
|
|
"boxes": boxes,
|
|
}, ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |