Spaces:
Sleeping
Sleeping
| """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() | |