Spaces:
Sleeping
Sleeping
File size: 1,586 Bytes
8bbb872 | 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 | """Run the eye open/closed model on one or more images."""
import sys
from pathlib import Path
from ultralytics import YOLO
def main():
project_root = Path(__file__).resolve().parent.parent
weight_candidates = [
project_root / "weights" / "best.pt",
project_root / "runs" / "classify" / "runs_cls" / "eye_open_closed_cpu" / "weights" / "best.pt",
project_root / "runs" / "classify" / "runs_cls" / "eye_open_closed_cpu" / "weights" / "last.pt",
]
weights = next((p for p in weight_candidates if p.is_file()), None)
if weights is None:
print("Weights not found. Put best.pt in weights/ or runs/.../weights/ (from model team).")
sys.exit(1)
if len(sys.argv) < 2:
print("Usage: python scripts/predict_image.py <image1> [image2 ...]")
print("Example: python scripts/predict_image.py path/to/image.png")
sys.exit(0)
model = YOLO(str(weights))
names = model.names
for path in sys.argv[1:]:
p = Path(path)
if not p.is_file():
print(p, "- file not found")
continue
try:
results = model.predict(str(p), imgsz=224, device="cpu", verbose=False)
except Exception as e:
print(p, "- error:", e)
continue
if not results:
print(p, "- no result")
continue
r = results[0]
top_idx = int(r.probs.top1)
conf = float(r.probs.top1conf)
label = names[top_idx]
print(f"{p.name}: {label} ({conf:.2%})")
if __name__ == "__main__":
main()
|