Update app.py

app.py

@@ -4,7 +4,7 @@ import zipfile
 import warnings
 warnings.filterwarnings('ignore')
 
-import streamlit as st
+import gradio as gr
 import torch
 import cv2
 import numpy as np
@@ -15,23 +15,30 @@ from torchvision.models.detection import fasterrcnn_resnet50_fpn
 from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
 from pathlib import Path
 
-@st.cache_resource
+# Global model variables
+MODEL = None
+DEVICE = None
+
 def load_model(model_path="seed_frcnn.pth", num_classes=2):
     """Load the trained Faster R-CNN model"""
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    global MODEL, DEVICE
+    
+    if MODEL is not None:
+        return MODEL, DEVICE
+    
+    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
     
     try:
-        model = fasterrcnn_resnet50_fpn(weights=None)
-        model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(
-            model.roi_heads.box_predictor.cls_score.in_features, num_classes
+        MODEL = fasterrcnn_resnet50_fpn(weights=None)
+        MODEL.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(
+            MODEL.roi_heads.box_predictor.cls_score.in_features, num_classes
         )
-        model.load_state_dict(torch.load(model_path, map_location=device))
-        model.to(device)
-        model.eval()
-        return model, device
+        MODEL.load_state_dict(torch.load(model_path, map_location=DEVICE))
+        MODEL.to(DEVICE)
+        MODEL.eval()
+        return MODEL, DEVICE
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None, None
+        raise Exception(f"Error loading model: {str(e)}")
 
 def run_inference(model, device, img_pil, score_thresh=0.5):
     """Run inference on a PIL image"""
@@ -69,19 +76,24 @@ def draw_boxes(image, detections):
     
     return cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB)
 
-def process_images(model, device, images, threshold, folder_name):
+def process_images(images, threshold, folder_name):
     """Process uploaded images and return crops as a ZIP file"""
+    if not images:
+        return None, "❌ No images uploaded!", None, None, None
+    
+    model, device = load_model()
+    
     if not folder_name:
         folder_name = "seed_crops"
     
     folder_name = "".join(c for c in folder_name if c.isalnum() or c in ('-', '_'))
     
-    zip_buffer = io.BytesIO()
+    zip_path = f"{folder_name}.zip"
     total_crops = 0
     processed_images = 0
     preview_images = []
     
-    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_file:
         for idx, img_file in enumerate(images):
             try:
                 img = Image.open(img_file).convert("RGB")
@@ -116,88 +128,94 @@ def process_images(model, device, images, threshold, folder_name):
                     processed_images += 1
                 
             except Exception as e:
-                st.error(f"Error processing image {idx}: {str(e)}")
+                print(f"Error processing image {idx}: {str(e)}")
                 continue
     
-    zip_buffer.seek(0)
-    return zip_buffer, total_crops, processed_images, preview_images
-
-# Streamlit UI
-st.set_page_config(page_title="Seed Detection & Cropping", page_icon="🌱", layout="wide")
-
-st.title("🌱 Seed Detection & Cropping Tool")
-st.markdown("Upload images to detect seeds using AI-powered Faster R-CNN model. Get all detected seeds as individual cropped images in a ZIP file.")
-
-# Load model
-with st.spinner("Loading model..."):
-    model, device = load_model("seed_frcnn.pth")
-
-if model is None:
-    st.error("Failed to load model. Please check if seed_frcnn.pth exists.")
-    st.stop()
-
-st.success(f"✅ Model loaded successfully on {device}!")
-
-# Sidebar controls
-st.sidebar.header("⚙️ Settings")
-threshold = st.sidebar.slider("🎚️ Detection Threshold", min_value=0.1, max_value=0.95, value=0.5, step=0.05)
-folder_name = st.sidebar.text_input("📁 Output Folder Name", value="seed_crops")
+    if total_crops > 0:
+        status_msg = f"✅ Processed {processed_images} images\n🌱 Detected and saved {total_crops} seed crops"
+        preview1 = preview_images[0] if len(preview_images) > 0 else None
+        preview2 = preview_images[1] if len(preview_images) > 1 else None
+        preview3 = preview_images[2] if len(preview_images) > 2 else None
+        return zip_path, status_msg, preview1, preview2, preview3
+    else:
+        return None, "⚠️ No seeds detected! Try lowering the threshold.", None, None, None
 
-st.sidebar.markdown("""
-### 💡 Tips:
-- **Lower threshold (0.3-0.5)**: More detections, may include false positives
-- **Higher threshold (0.6-0.8)**: Conservative, high-confidence only
-""")
+# Initialize model on startup
+try:
+    load_model()
+    model_status = "✅ Model loaded successfully!"
+except Exception as e:
+    model_status = f"❌ Error loading model: {str(e)}"
 
-# File uploader
-uploaded_files = st.file_uploader("📤 Upload Images", type=["jpg", "jpeg", "png", "bmp"], accept_multiple_files=True)
-
-if uploaded_files:
-    st.info(f"📷 {len(uploaded_files)} image(s) uploaded")
+# Create Gradio interface
+with gr.Blocks(title="Seed Detection & Cropping", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌱 Seed Detection & Cropping Tool")
+    gr.Markdown("Upload images to detect seeds using AI-powered Faster R-CNN model. Get all detected seeds as individual cropped images in a ZIP file.")
+    gr.Markdown(f"**Model Status:** {model_status}")
     
-    if st.button("🔍 Detect & Crop Seeds", type="primary"):
-        with st.spinner("Processing images..."):
-            zip_buffer, total_crops, processed_images, preview_images = process_images(
-                model, device, uploaded_files, threshold, folder_name
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("## ⚙️ Settings")
+            threshold = gr.Slider(
+                minimum=0.1, 
+                maximum=0.95, 
+                value=0.5, 
+                step=0.05, 
+                label="🎚️ Detection Threshold",
+                info="Lower = more detections, Higher = only high-confidence"
+            )
+            folder_name = gr.Textbox(
+                value="seed_crops", 
+                label="📁 Output Folder Name",
+                placeholder="seed_crops"
             )
             
-            if total_crops > 0:
-                st.success(f"✅ Processed {processed_images} images")
-                st.success(f"🌱 Detected and saved {total_crops} seed crops")
-                
-                # Download button
-                st.download_button(
-                    label="💾 Download ZIP File",
-                    data=zip_buffer.getvalue(),
-                    file_name=f"{folder_name}.zip",
-                    mime="application/zip"
-                )
-                
-                # Preview
-                if preview_images:
-                    st.markdown("### 🖼️ Preview (First 3 images with detections)")
-                    cols = st.columns(min(3, len(preview_images)))
-                    for idx, preview_img in enumerate(preview_images):
-                        cols[idx].image(preview_img, use_column_width=True)
-            else:
-                st.warning("⚠️ No seeds detected! Try lowering the threshold.")
-else:
-    st.info("👆 Upload images to get started")
+            gr.Markdown("""
+            ### 💡 Tips:
+            - **Lower threshold (0.3-0.5)**: More detections, may include false positives
+            - **Higher threshold (0.6-0.8)**: Conservative, high-confidence only
+            """)
+        
+        with gr.Column(scale=2):
+            images = gr.File(
+                file_count="multiple", 
+                label="📤 Upload Images",
+                file_types=["image"]
+            )
+            
+            process_btn = gr.Button("🔍 Detect & Crop Seeds", variant="primary", size="lg")
+            
+            status = gr.Textbox(label="Status", interactive=False)
+            download = gr.File(label="💾 Download ZIP File")
+    
+    gr.Markdown("### 🖼️ Preview (First 3 images with detections)")
+    with gr.Row():
+        preview1 = gr.Image(label="Preview 1")
+        preview2 = gr.Image(label="Preview 2")
+        preview3 = gr.Image(label="Preview 3")
+    
+    with gr.Accordion("📋 How to Use", open=False):
+        gr.Markdown("""
+        1. **Upload Images**: Click the upload box and select one or multiple images
+        2. **Adjust Threshold**: Use the slider to control detection sensitivity
+        3. **Set Folder Name**: Enter the folder name for your crops
+        4. **Click Detect**: Process your images
+        5. **Download**: Get your ZIP file with all cropped seeds!
+        
+        ### Output Format:
+        ```
+        your_folder_name/
+        ├── image1_seed_000_score_0.850.jpg
+        ├── image1_seed_001_score_0.720.jpg
+        └── image2_seed_000_score_0.910.jpg
+        ```
+        """)
+    
+    process_btn.click(
+        fn=process_images,
+        inputs=[images, threshold, folder_name],
+        outputs=[download, status, preview1, preview2, preview3]
+    )
 
-# Instructions
-with st.expander("📋 How to Use"):
-    st.markdown("""
-    1. **Upload Images**: Click the upload box and select one or multiple images
-    2. **Adjust Threshold**: Use the slider in the sidebar to control detection sensitivity
-    3. **Set Folder Name**: Enter the folder name for your crops
-    4. **Click Detect**: Process your images
-    5. **Download**: Get your ZIP file with all cropped seeds!
-    
-    ### Output Format:
-    ```
-    your_folder_name/
-    ├── image1_seed_000_score_0.850.jpg
-    ├── image1_seed_001_score_0.720.jpg
-    └── image2_seed_000_score_0.910.jpg
-    ```
-    """)
+if __name__ == "__main__":
+    demo.launch()