169 lines
4.5 KiB
Python
169 lines
4.5 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 existing_file(path):
|
|
try:
|
|
p = Path(path).expanduser().resolve()
|
|
if p.exists() and p.is_file() and p.stat().st_size > 0:
|
|
return p
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def resolve_model_path(root, requested):
|
|
if requested:
|
|
p = existing_file(requested)
|
|
if p:
|
|
return p, "yolo26_model"
|
|
return Path(requested).expanduser(), "detector_missing"
|
|
|
|
trained = root / "detector" / "model" / "best.pt"
|
|
p = existing_file(trained)
|
|
if p:
|
|
return p, "yolo26_detector"
|
|
|
|
return trained, "detector_missing"
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
parser.add_argument("--image", required=True)
|
|
parser.add_argument("--model", default="")
|
|
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, model_source = resolve_model_path(root, args.model)
|
|
if not model_path.exists():
|
|
print(json.dumps({
|
|
"available": False,
|
|
"source": model_source,
|
|
"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"
|
|
|
|
try:
|
|
results = model.predict(
|
|
source=str(image_path),
|
|
conf=float(args.conf),
|
|
imgsz=int(args.imgsz),
|
|
verbose=False,
|
|
device=device,
|
|
)
|
|
except Exception as e:
|
|
print(json.dumps({
|
|
"available": False,
|
|
"source": "detector_predict_failed",
|
|
"modelPath": str(model_path),
|
|
"image": str(image_path),
|
|
"error": repr(e),
|
|
"boxes": [],
|
|
}, ensure_ascii=False))
|
|
return
|
|
|
|
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()]
|
|
|
|
x1 = max(0.0, min(float(img_w), x1))
|
|
y1 = max(0.0, min(float(img_h), y1))
|
|
x2 = max(0.0, min(float(img_w), x2))
|
|
y2 = max(0.0, min(float(img_h), y2))
|
|
|
|
if x2 <= x1 or y2 <= y1:
|
|
continue
|
|
|
|
x = x1 / img_w
|
|
y = y1 / img_h
|
|
w = (x2 - x1) / img_w
|
|
h = (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": model_source,
|
|
"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()
|