Spaces:
Sleeping
Sleeping
| """Gradio interface for the Analogic Watch Detector model. | |
| This module exposes a lightweight Gradio demo that can be used on | |
| Hugging Face Spaces. It loads the YOLO model once, runs inference on the | |
| uploaded image and renders the predicted time alongside the annotated | |
| image. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| from typing import Optional, Tuple | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from utils.clock_utils import process_clock_time | |
| from utils.detections_utils import get_latest_train_dir, run_detection | |
| _MODEL: Optional[YOLO] = None | |
| def _resolve_model_path() -> str: | |
| """Return the best available model path.""" | |
| env_path = os.environ.get("MODEL_PATH") | |
| if env_path and os.path.exists(env_path): | |
| return env_path | |
| # Priority 1: Specific tuned model for HF deployment | |
| tune4_model = "tune4_best.pt" | |
| if os.path.exists(tune4_model): | |
| return tune4_model | |
| default_weight = "yolov8s.pt" | |
| if os.path.exists(default_weight): | |
| return default_weight | |
| try: | |
| return os.path.join(get_latest_train_dir(), "weights", "best.pt") | |
| except FileNotFoundError as exc: # pragma: no cover - defensive path | |
| raise RuntimeError( | |
| "Model weights were not found. Provide them via the MODEL_PATH " | |
| "environment variable or include 'yolov8s.pt' in the repository." | |
| ) from exc | |
| def _load_model() -> YOLO: | |
| """Lazy-load the YOLO model to keep the interface responsive.""" | |
| global _MODEL | |
| if _MODEL is None: | |
| model_path = _resolve_model_path() | |
| _MODEL = YOLO(model_path) | |
| return _MODEL | |
| def _format_time(prediction: Optional[dict]) -> str: | |
| """Generate a human readable string for the detected time.""" | |
| if not prediction: | |
| return "Unable to determine the time from the detected clock." | |
| hours = prediction.get("hours") | |
| minutes = prediction.get("minutes") | |
| seconds = prediction.get("seconds") | |
| if hours is None: | |
| return "Unable to determine the time from the detected clock." | |
| if minutes is None: | |
| return f"Detected hour hand at {hours:02d}." | |
| if seconds is None: | |
| return f"Detected time: {hours:02d}:{minutes:02d}." | |
| return f"Detected time: {hours:02d}:{minutes:02d}:{seconds:02d}." | |
| def predict(image: np.ndarray, confidence: float) -> Tuple[np.ndarray, str]: | |
| """Run detection on the uploaded image and return the annotated preview.""" | |
| if image is None: | |
| raise gr.Error("Please upload an image of an analog clock.") | |
| # Basic guard against oversized inputs to reduce DoS risk | |
| try: | |
| if hasattr(image, "nbytes") and (image.nbytes > 40 * 1024 * 1024): | |
| raise gr.Error("Image is too large. Please upload a smaller file.") | |
| if image.size and image.size > 20_000_000: | |
| raise gr.Error("Image resolution is too high. Please downscale and retry.") | |
| except Exception: | |
| pass | |
| image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) | |
| detections, results = run_detection( | |
| image=image_bgr, | |
| image_path=None, | |
| confidence=confidence, | |
| save_path=None, | |
| save_visualization=False, | |
| return_prediction_results=True, | |
| model=_load_model(), | |
| ) | |
| if not detections or not detections[0]: | |
| return image, "No clock components detected in the provided image." | |
| prediction = process_clock_time(detections, "uploaded_image") | |
| annotated = image | |
| if results: | |
| annotated_bgr = results[0].plot() | |
| annotated = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB) | |
| return annotated, _format_time(prediction) | |
| def build_interface() -> gr.Blocks: | |
| """Create the Gradio Blocks interface.""" | |
| with gr.Blocks(title="Analog Clock Time Detector") as demo: | |
| gr.Markdown( | |
| """ | |
| # Analog Clock Time Detector | |
| Upload a picture of an analog clock to detect the time displayed on it. | |
| The model is based on YOLOv8 and predicts the hour, minute and second | |
| hands when available. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_input = gr.Image( | |
| type="numpy", | |
| label="Clock image", | |
| image_mode="RGB", | |
| height=600, | |
| ) | |
| confidence_slider = gr.Slider( | |
| minimum=0.01, | |
| maximum=0.5, | |
| step=0.01, | |
| value=0.1, | |
| label="Detection confidence threshold", | |
| ) | |
| submit_btn = gr.Button("Detect time") | |
| with gr.Column(): | |
| annotated_image = gr.Image( | |
| type="numpy", | |
| label="Detections", | |
| ) | |
| time_output = gr.Textbox( | |
| label="Predicted time", | |
| placeholder="The predicted time will appear here.", | |
| ) | |
| submit_btn.click( | |
| fn=predict, | |
| inputs=[image_input, confidence_slider], | |
| outputs=[annotated_image, time_output], | |
| ) | |
| # Load examples from the examples directory | |
| example_images = [] | |
| if os.path.exists("examples"): | |
| example_images = [ | |
| os.path.join("examples", f) | |
| for f in os.listdir("examples") | |
| if f.lower().endswith(('.png', '.jpg', '.jpeg')) | |
| ] | |
| # Sort for consistent order, though not strictly required | |
| example_images.sort() | |
| # Fallback to img directory if no examples found (local dev fallback/legacy) | |
| if not example_images and os.path.exists("img"): | |
| example_images = [ | |
| os.path.join("img", "1.png"), | |
| os.path.join("img", "2.png"), | |
| ] | |
| if example_images: | |
| gr.Examples( | |
| examples=[[img, 0.1] for img in example_images], | |
| inputs=[image_input, confidence_slider], | |
| outputs=[annotated_image, time_output], | |
| fn=predict, | |
| cache_examples=False, | |
| ) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = build_interface() | |
| # Detect Hugging Face Spaces environment | |
| on_hf = bool(os.environ.get("SPACE_ID")) | |
| # Configure queue | |
| demo.queue(default_concurrency_limit=1, max_size=8) | |
| if on_hf: | |
| # Let Spaces manage networking/binding | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |
| else: | |
| # Local dev: bind only to localhost and avoid public share | |
| try: | |
| demo.launch(server_name="127.0.0.1", share=False, allowed_paths=["img"]) | |
| except TypeError: | |
| demo.launch(server_name="127.0.0.1", share=False) | |