Spaces:
Sleeping
Sleeping
| from typing import Tuple, Dict | |
| import gradio as gr | |
| import supervision as sv | |
| import numpy as np | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download | |
| from ultralytics import YOLO | |
| import cv2 | |
| # Define models | |
| MODEL_OPTIONS = { | |
| "YOLOv11-Nano": "yolov11n-seg.pt", | |
| "YOLOv11-Small": "yolov11s-seg.pt", | |
| "YOLOv11-Medium": "yolov11m-seg.pt", | |
| "YOLOv11-Large": "yolov11l-seg.pt", | |
| "YOLOv11-XLarge": "yolov11x-seg.pt" | |
| } | |
| print(MODEL_OPTIONS) | |
| # Dictionary to store loaded models | |
| models: Dict[str, YOLO] = {} | |
| # Load all models | |
| for name, model_file in MODEL_OPTIONS.items(): | |
| model_path = hf_hub_download( | |
| repo_id="wjbmattingly/kraken-yiddish", | |
| filename=model_file | |
| ) | |
| models[name] = YOLO(model_path) | |
| # Create annotators | |
| LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK) | |
| MASK_ANNOTATOR = sv.MaskAnnotator() | |
| BOX_ANNOTATOR = sv.BoxAnnotator() | |
| def detect_and_annotate( | |
| image: np.ndarray, | |
| model_name: str, | |
| conf_threshold: float, | |
| iou_threshold: float | |
| ) -> np.ndarray: | |
| # Get the selected model | |
| model = models[model_name] | |
| try: | |
| # Perform inference with YOLO standard settings | |
| results = model( | |
| image, | |
| conf=conf_threshold, | |
| iou=iou_threshold, | |
| verbose=False, | |
| device='cpu' | |
| ) | |
| # Get the first result and create annotated image | |
| result = results[0] # Get first result from list | |
| annotated_image = result.plot( | |
| conf=True, # Show confidence scores | |
| line_width=None, # Auto-scale line width | |
| font_size=None, # Auto-scale font size | |
| boxes=True, # Draw boxes | |
| masks=True, # Draw masks if available | |
| probs=True, # Draw classification probabilities if available | |
| labels=True # Draw labels | |
| ) | |
| return annotated_image | |
| except Exception as e: | |
| print(f"\nError during detection: {str(e)}") | |
| import traceback | |
| traceback.print_exc() | |
| return image.copy() | |
| # Create Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Yiddish Document Line Parser") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image( | |
| label="Input Image", | |
| type='numpy' | |
| ) | |
| with gr.Accordion("Detection Settings", open=True): | |
| model_selector = gr.Dropdown( | |
| choices=list(MODEL_OPTIONS.keys()), | |
| value=list(MODEL_OPTIONS.keys())[0], | |
| label="Model", | |
| info="Select YOLO model variant" | |
| ) | |
| with gr.Row(): | |
| conf_threshold = gr.Slider( | |
| label="Confidence Threshold", | |
| minimum=0.0, | |
| maximum=1.0, | |
| step=0.05, | |
| value=0.25, | |
| ) | |
| iou_threshold = gr.Slider( | |
| label="IoU Threshold", | |
| minimum=0.0, | |
| maximum=1.0, | |
| step=0.05, | |
| value=0.45, | |
| info="Decrease for stricter detection, increase for more overlapping boxes" | |
| ) | |
| with gr.Row(): | |
| clear_btn = gr.Button("Clear") | |
| detect_btn = gr.Button("Detect", variant="primary") | |
| with gr.Column(): | |
| output_image = gr.Image( | |
| label="Detection Result", | |
| type='numpy' | |
| ) | |
| def process_image( | |
| image: np.ndarray, | |
| model_name: str, | |
| conf_threshold: float, | |
| iou_threshold: float | |
| ) -> Tuple[np.ndarray, np.ndarray]: | |
| if image is None: | |
| return None, None | |
| annotated_image = detect_and_annotate(image, model_name, conf_threshold, iou_threshold) | |
| return image, annotated_image | |
| def clear(): | |
| return None, None | |
| # Connect buttons to functions | |
| detect_btn.click( | |
| process_image, | |
| inputs=[input_image, model_selector, conf_threshold, iou_threshold], | |
| outputs=[input_image, output_image] | |
| ) | |
| clear_btn.click( | |
| clear, | |
| inputs=None, | |
| outputs=[input_image, output_image] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(debug=True, show_error=True) |