Spaces:
Running
Running
| import gradio as gr | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| import os | |
| # 1. Load Models | |
| models = { | |
| "CowCatcherV15": YOLO('cowcatcherV15.pt'), | |
| "CowCatcherV16 (Experimental)": YOLO('cowcatcherV16.pt') | |
| } | |
| # 2. Styling & Colors | |
| custom_css = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@400;700&display=swap'); | |
| @import url('https://fonts.googleapis.com/css2?family=Bebas+Neue&display=swap'); | |
| .bebas-font { | |
| font-family: 'Bebas Neue', sans-serif !important; | |
| text-transform: uppercase; | |
| letter-spacing: 1px; | |
| } | |
| .roboto-font { | |
| font-family: 'Roboto', sans-serif !important; | |
| } | |
| .primary-btn { | |
| background-color: #386938 !important; | |
| border: none !important; | |
| font-family: 'Bebas Neue', sans-serif !important; | |
| font-size: 1.2rem !important; | |
| color: white !important; | |
| } | |
| .gradio-container label span { | |
| font-family: 'Roboto', sans-serif !important; | |
| font-weight: bold; | |
| } | |
| """ | |
| description_text = """ | |
| **CowCatcher AI** is an open-source computer vision model designed to monitor your herd 24/7. | |
| By analyzing footage from your barn cameras, it automatically detects "mounting" behavior — the primary sign of estrus (heat). | |
| > ⚠️ **Best Results: High-Angle View** — Trained for security camera perspectives (4–5 meters high). Eye-level photos may not work well. | |
| """ | |
| # 3. Prediction function | |
| def predict(img, model_name, conf_threshold): | |
| if img is None: | |
| return None | |
| results = models[model_name](img, conf=conf_threshold) | |
| res_plotted = results[0].plot() | |
| return Image.fromarray(res_plotted[:, :, ::-1]) | |
| # 4. Load example images | |
| example_folder = "examples" | |
| example_images = sorted([ | |
| os.path.join(example_folder, f) | |
| for f in os.listdir(example_folder) | |
| if f.lower().endswith(('.png', '.jpg', '.jpeg')) | |
| ]) if os.path.exists(example_folder) else [] | |
| # 5. Build UI | |
| with gr.Blocks(title="CowCatcher AI") as demo: | |
| with gr.Column(elem_classes="roboto-font"): | |
| gr.Markdown("# 🐮 CowCatcher AI", elem_classes="bebas-font") | |
| gr.Markdown(description_text, elem_classes="roboto-font") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(type="pil", label="Upload Barn Image") | |
| with gr.Row(): | |
| model_drop = gr.Dropdown( | |
| choices=list(models.keys()), | |
| value="CowCatcherV15", | |
| label="Select Model" | |
| ) | |
| conf_slider = gr.Slider( | |
| minimum=0.0, | |
| maximum=1.0, | |
| value=0.7, | |
| step=0.05, | |
| label="Confidence Threshold", | |
| info="Higher = fewer but more certain detections." | |
| ) | |
| predict_btn = gr.Button("🔍 Run Detection", variant="primary", elem_classes="primary-btn") | |
| with gr.Column(): | |
| output_img = gr.Image(type="pil", label="Detection Result") | |
| if example_images: | |
| gr.Markdown("### 📸 Try an Example", elem_classes="bebas-font") | |
| gr.Examples( | |
| examples=example_images, | |
| inputs=input_img, | |
| label=None | |
| ) | |
| predict_btn.click( | |
| fn=predict, | |
| inputs=[input_img, model_drop, conf_slider], | |
| outputs=output_img | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch( | |
| css=custom_css, | |
| theme=gr.themes.Soft(primary_hue="green"), | |
| ) |