Update app.py

app.py

@@ -7,305 +7,245 @@ from ultralytics import YOLO
 import supervision as sv
 from PIL import Image
 from huggingface_hub import snapshot_download
-from functools import lru_cache
-from typing import Tuple, Optional
 import spaces
+from typing import Tuple, Optional
+import logging
 
-# Constants
-DEFAULT_CONFIDENCE_THRESHOLD = 0.1
-DEFAULT_NMS_THRESHOLD = 0.0
-DEFAULT_SLICE_WIDTH = 1024
-DEFAULT_SLICE_HEIGHT = 1024
-DEFAULT_OVERLAP_WIDTH = 0
-DEFAULT_OVERLAP_HEIGHT = 0
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Detection parameters
+CONFIDENCE_THRESHOLD = 0.1
+NMS_THRESHOLD = 0.0
+SLICE_WIDTH = 1024
+SLICE_HEIGHT = 1024
+OVERLAP_WIDTH = 0
+OVERLAP_HEIGHT = 0
 ANNOTATION_COLOR = sv.Color.RED
 ANNOTATION_THICKNESS = 4
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-REPO_ID = 'edeler/ICC'
-
-# Global model cache
-_model_cache = None
-
-@lru_cache(maxsize=1)
-def download_model() -> str:
-    """Download and cache model from Hugging Face Hub."""
-    model_dir = snapshot_download(REPO_ID, local_dir='./models/ICC')
-    return os.path.join(model_dir, "best.pt")
-
-def load_model() -> YOLO:
-    """Lazy-load and cache the YOLO model."""
-    global _model_cache
-    if _model_cache is None:
-        model_path = download_model()
-        _model_cache = YOLO(model_path).to(DEVICE)
-    return _model_cache
-
-def validate_image(image: Optional[np.ndarray]) -> bool:
-    """Validate input image."""
-    if image is None:
-        return False
-    if not isinstance(image, np.ndarray):
-        return False
-    if image.size == 0:
-        return False
-    return True
 
-def preprocess_image(image: np.ndarray) -> np.ndarray:
-    """Convert image to BGR format if needed."""
-    # Gradio provides RGB images, convert to BGR for OpenCV/YOLO
-    if len(image.shape) == 3 and image.shape[-1] == 3:
-        # Check if already BGR or RGB by inspecting typical color distributions
-        return cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    return image
+# Device configuration
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+logger.info(f"Using device: {device}")
 
-def create_detection_callback(model: YOLO, confidence_threshold: float):
-    """Factory function to create detection callback with closure over parameters."""
-    def callback(image_slice: np.ndarray) -> sv.Detections:
-        with torch.no_grad():  # Disable gradient computation for inference
-            result = model(image_slice, verbose=False)[0]
-        detections = sv.Detections.from_ultralytics(result)
-        return detections[detections.confidence >= confidence_threshold]
-    return callback
-
-def perform_detection(
-    image: np.ndarray,
-    confidence_threshold: float,
-    nms_threshold: float,
-    slice_width: int,
-    slice_height: int,
-    overlap_width: int,
-    overlap_height: int
-) -> Tuple[sv.Detections, int]:
-    """Perform object detection with slicing and NMS."""
-    model = load_model()
-    
-    # Create slicer with callback
-    callback = create_detection_callback(model, confidence_threshold)
-    slicer = sv.InferenceSlicer(
-        callback=callback,
-        slice_wh=(slice_width, slice_height),
-        overlap_wh=(overlap_width, overlap_height),
-        overlap_ratio_wh=None
-    )
-    
-    # Perform slicing-based inference
-    detections = slicer(image)
-    
-    # Apply NMS
-    if nms_threshold > 0:
-        detections = detections.with_nms(threshold=nms_threshold, class_agnostic=False)
-    
-    return detections, len(detections)
+# Model initialization
+def load_model():
+    """Load YOLO model from Hugging Face Hub with error handling."""
+    try:
+        repo_id = 'edeler/ICC'
+        logger.info(f"Downloading model from {repo_id}...")
+        model_dir = snapshot_download(repo_id, local_dir='./models/ICC')
+        model_path = os.path.join(model_dir, "best.pt")
+        
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"Model file not found at {model_path}")
+        
+        logger.info(f"Loading model from {model_path}...")
+        model = YOLO(model_path)
+        model.to(device)
+        logger.info("Model loaded successfully")
+        return model
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        raise
 
-def annotate_image(image: np.ndarray, detections: sv.Detections) -> np.ndarray:
-    """Annotate image with detection bounding boxes."""
-    box_annotator = sv.OrientedBoxAnnotator(
-        color=ANNOTATION_COLOR,
-        thickness=ANNOTATION_THICKNESS
-    )
-    return box_annotator.annotate(scene=image.copy(), detections=detections)
+# Load model once at startup
+model = load_model()
 
-@spaces.GPU
-def detect_objects(
-    image: Optional[np.ndarray],
-    confidence_threshold: float,
-    nms_threshold: float,
-    slice_width: int,
-    slice_height: int,
-    overlap_width: int,
-    overlap_height: int
-) -> Tuple[Optional[Image.Image], str]:
+@spaces.GPU(duration=60)  # Allocate GPU for up to 60 seconds
+def detect_objects(image: Optional[np.ndarray]) -> Tuple[Optional[Image.Image], str]:
     """
-    Main detection function for Gradio interface.
+    Detect objects in the input image using YOLO model with sliced inference.
     
     Args:
-        image: Input image as numpy array
-        confidence_threshold: Minimum confidence for detections
-        nms_threshold: IoU threshold for non-maximum suppression
-        slice_width: Width of each detection slice
-        slice_height: Height of each detection slice
-        overlap_width: Overlap width between slices
-        overlap_height: Overlap height between slices
-    
+        image: Input image as numpy array (RGB format from Gradio)
+        
     Returns:
-        Tuple of (annotated image, detection summary)
+        Tuple of (annotated PIL Image, detection summary string)
     """
+    # Validate input
+    if image is None:
+        return None, "⚠️ Please upload an image first."
+    
     try:
-        # Validate input
-        if not validate_image(image):
-            return None, "⚠️ Please upload a valid image."
+        # Convert RGB (from Gradio) to BGR (for OpenCV/YOLO)
+        image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
         
-        # Preprocess image
-        image_bgr = preprocess_image(image)
+        # Define callback for sliced inference
+        def inference_callback(image_slice: np.ndarray) -> sv.Detections:
+            """Process each image slice."""
+            result = model(image_slice, verbose=False)[0]
+            detections = sv.Detections.from_ultralytics(result)
+            # Filter by confidence threshold
+            return detections[detections.confidence >= CONFIDENCE_THRESHOLD]
+        
+        # Initialize slicer
+        slicer = sv.InferenceSlicer(
+            callback=inference_callback,
+            slice_wh=(SLICE_WIDTH, SLICE_HEIGHT),
+            overlap_wh=(OVERLAP_WIDTH, OVERLAP_HEIGHT),
+            overlap_ratio_wh=None
+        )
+        
+        # Perform inference
+        logger.info("Running detection...")
+        detections = slicer(image_bgr)
         
-        # Perform detection
-        detections, total_count = perform_detection(
-            image_bgr,
-            confidence_threshold,
-            nms_threshold,
-            slice_width,
-            slice_height,
-            overlap_width,
-            overlap_height
+        # Apply NMS to remove duplicate detections
+        detections = detections.with_nms(
+            threshold=NMS_THRESHOLD,
+            class_agnostic=False
         )
         
+        # Count detections
+        total_detections = len(detections)
+        logger.info(f"Found {total_detections} detections")
+        
         # Annotate image
-        annotated_img = annotate_image(image_bgr, detections)
+        box_annotator = sv.OrientedBoxAnnotator(
+            color=ANNOTATION_COLOR,
+            thickness=ANNOTATION_THICKNESS
+        )
+        annotated_img_bgr = box_annotator.annotate(
+            scene=image_bgr.copy(),
+            detections=detections
+        )
         
         # Convert back to RGB for display
-        annotated_img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
+        annotated_img_rgb = cv2.cvtColor(annotated_img_bgr, cv2.COLOR_BGR2RGB)
         
-        # Create summary with additional info
-        summary = f"✅ Total Detections: {total_count}"
-        if total_count > 0:
-            avg_confidence = np.mean(detections.confidence) if len(detections.confidence) > 0 else 0
-            summary += f"\n📊 Average Confidence: {avg_confidence:.2%}"
+        # Create result message
+        result_msg = f"✅ Total Detections: {total_detections}"
         
-        return Image.fromarray(annotated_img_rgb), summary
+        return Image.fromarray(annotated_img_rgb), result_msg
         
     except Exception as e:
-        return None, f"❌ Error during detection: {str(e)}"
+        error_msg = f"❌ Error during detection: {str(e)}"
+        logger.error(error_msg)
+        return None, error_msg
 
 def get_example_images() -> list:
     """Get list of example images from the current directory."""
-    example_root = os.path.dirname(__file__)
-    return [
-        os.path.join(example_root, file)
-        for file in os.listdir(example_root)
-        if file.lower().endswith(('.jpg', '.jpeg', '.png'))
-    ]
+    try:
+        example_root = os.path.dirname(__file__) or "."
+        example_images = [
+            os.path.join(example_root, file)
+            for file in os.listdir(example_root)
+            if file.lower().endswith((".jpg", ".jpeg", ".png"))
+        ]
+        return example_images[:10]  # Limit to 10 examples
+    except Exception as e:
+        logger.warning(f"Could not load example images: {str(e)}")
+        return []
 
-def create_interface() -> gr.Blocks:
+def create_interface():
     """Create and configure the Gradio interface."""
-    with gr.Blocks(title="ICC Detection Tool") as demo:
+    
+    with gr.Blocks(
+        title="ICC Detection Tool",
+        theme=gr.themes.Soft(),
+        css="""
+        .gradio-container {max-width: 1200px !important}
+        #title {text-align: center; color: #2563eb;}
+        """
+    ) as demo:
+        
         gr.Markdown(
             """
             # 🔬 Interstitial Cell of Cajal Detection and Quantification Tool
             
-            Upload an image to detect and count Interstitial Cells of Cajal (ICC).
-            Adjust the parameters below for fine-tuned detection.
-            """
+            Upload an image to detect and quantify Interstitial Cells of Cajal (ICC).
+            The model uses advanced YOLO detection with sliced inference for accurate results.
+            """,
+            elem_id="title"
         )
         
         with gr.Row():
             with gr.Column(scale=1):
                 input_img = gr.Image(
-                    label="📤 Upload an Image",
+                    label="📤 Upload Image",
                     type="numpy",
-                    interactive=True
+                    height=400
                 )
                 
-                # Advanced settings in accordion
-                with gr.Accordion("⚙️ Advanced Settings", open=False):
-                    confidence_slider = gr.Slider(
-                        minimum=0.01,
-                        maximum=1.0,
-                        value=DEFAULT_CONFIDENCE_THRESHOLD,
-                        step=0.01,
-                        label="Confidence Threshold",
-                        info="Minimum confidence for detections"
-                    )
-                    nms_slider = gr.Slider(
-                        minimum=0.0,
-                        maximum=1.0,
-                        value=DEFAULT_NMS_THRESHOLD,
-                        step=0.05,
-                        label="NMS Threshold",
-                        info="IoU threshold for non-maximum suppression"
-                    )
-                    
-                    with gr.Row():
-                        slice_width = gr.Number(
-                            value=DEFAULT_SLICE_WIDTH,
-                            label="Slice Width",
-                            precision=0
-                        )
-                        slice_height = gr.Number(
-                            value=DEFAULT_SLICE_HEIGHT,
-                            label="Slice Height",
-                            precision=0
-                        )
-                    
-                    with gr.Row():
-                        overlap_width = gr.Number(
-                            value=DEFAULT_OVERLAP_WIDTH,
-                            label="Overlap Width",
-                            precision=0
-                        )
-                        overlap_height = gr.Number(
-                            value=DEFAULT_OVERLAP_HEIGHT,
-                            label="Overlap Height",
-                            precision=0
-                        )
-                
                 with gr.Row():
-                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
-                    detect_btn = gr.Button("🔍 Detect", variant="primary")
+                    clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
+                    detect_btn = gr.Button("🔍 Detect", variant="primary", scale=2)
+                
+                # Example images
+                example_images = get_example_images()
+                if example_images:
+                    with gr.Accordion("📁 Example Images", open=False):
+                        gr.Examples(
+                            examples=[[img] for img in example_images],
+                            inputs=[input_img],
+                            label=None
+                        )
             
             with gr.Column(scale=1):
                 output_img = gr.Image(
-                    label="📊 Detection Result",
-                    interactive=False
+                    label="✨ Detection Result",
+                    type="pil",
+                    height=400
                 )
                 detection_count = gr.Textbox(
-                    label="📈 Detection Summary",
+                    label="📊 Detection Summary",
                     interactive=False,
-                    lines=3
+                    lines=2
                 )
         
-        # Examples section
-        example_images = get_example_images()
-        if example_images:
-            with gr.Accordion("📁 Example Images", open=False):
-                gr.Examples(
-                    examples=[[img] for img in example_images],
-                    inputs=[input_img],
-                    label="Click to load"
-                )
+        # Model information
+        with gr.Accordion("ℹ️ Model Information", open=False):
+            gr.Markdown(
+                f"""
+                **Configuration:**
+                - Confidence Threshold: {CONFIDENCE_THRESHOLD}
+                - NMS Threshold: {NMS_THRESHOLD}
+                - Slice Size: {SLICE_WIDTH}x{SLICE_HEIGHT}
+                - Device: {device.upper()}
+                - Model: edeler/ICC
+                """
+            )
         
-        # Button actions
-        detect_btn.click(
-            fn=detect_objects,
-            inputs=[
-                input_img,
-                confidence_slider,
-                nms_slider,
-                slice_width,
-                slice_height,
-                overlap_width,
-                overlap_height
-            ],
-            outputs=[output_img, detection_count]
-        )
+        # Event handlers
+        def reset_interface():
+            """Reset all components to initial state."""
+            return None, None, ""
         
         clear_btn.click(
-            fn=lambda: (None, None, ""),
+            fn=reset_interface,
             inputs=None,
             outputs=[input_img, output_img, detection_count]
         )
         
-        gr.Markdown(
-            """
-            ---
-            ### 💡 Tips:
-            - **Confidence Threshold**: Lower values detect more objects but may include false positives
-            - **NMS Threshold**: Higher values keep more overlapping boxes
-            - **Slice Settings**: Adjust for large images or when objects are at different scales
-            """
+        detect_btn.click(
+            fn=detect_objects,
+            inputs=[input_img],
+            outputs=[output_img, detection_count],
+            api_name="detect"
+        )
+        
+        # Allow Enter key to trigger detection
+        input_img.upload(
+            fn=lambda: "⏳ Image uploaded. Click 'Detect' to start...",
+            inputs=None,
+            outputs=detection_count
         )
     
     return demo
 
+# Main execution
 if __name__ == "__main__":
-    # Preload model on startup for faster first inference
-    print(f"🚀 Loading model on {DEVICE}...")
-    load_model()
-    print("✅ Model loaded successfully!")
-    
-    # Create and launch interface
-    demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    try:
+        demo = create_interface()
+        demo.queue(max_size=20)  # Enable queuing for better handling of concurrent requests
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,  # Set to True if you want a public link
+            show_error=True
+        )
+    except Exception as e:
+        logger.error(f"Failed to launch app: {str(e)}")
+        raise