Update app.py

app.py

@@ -4,10 +4,10 @@ import logging
 import os
 import sys
 import threading
+import traceback
 from collections import Counter
 from io import BytesIO
 from typing import Dict, List, Optional, Tuple, Union
-
 import gradio as gr
 import pandas as pd
 import requests
@@ -22,17 +22,19 @@ from transformers import (
     DetrImageProcessor,
     YolosForObjectDetection,
     YolosImageProcessor,
+    DeformableDetrForObjectDetection,
+    DeformableDetrImageProcessor,
 )
 import nest_asyncio
-
 # ------------------------------
 # Configuration
 # ------------------------------
-
 # Configure logging for debugging and monitoring
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
-
+# Define device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger.info(f"Using device: {device}")
 # Define constants for model and server configuration
 CONFIDENCE_THRESHOLD: float = 0.5  # Default threshold for object detection confidence
 VALID_MODELS: List[str] = [
@@ -41,7 +43,9 @@ VALID_MODELS: List[str] = [
     "facebook/detr-resnet-50-panoptic",
     "facebook/detr-resnet-101-panoptic",
     "hustvl/yolos-tiny",
+    "hustvl/yolos-small",
     "hustvl/yolos-base",
+    "SenseTime/deformable-detr",
 ]
 MODEL_DESCRIPTIONS: Dict[str, str] = {
     "facebook/detr-resnet-50": "DETR with ResNet-50 for object detection. Fast and accurate.",
@@ -49,22 +53,21 @@ MODEL_DESCRIPTIONS: Dict[str, str] = {
     "facebook/detr-resnet-50-panoptic": "DETR with ResNet-50 for panoptic segmentation.",
     "facebook/detr-resnet-101-panoptic": "DETR with ResNet-101 for panoptic segmentation.",
     "hustvl/yolos-tiny": "YOLOS Tiny. Lightweight and fast.",
-    "hustvl/yolos-base": "YOLOS Base. Balances speed and accuracy."
+    "hustvl/yolos-small": "YOLOS Small. Medium speed and accuracy.",
+    "hustvl/yolos-base": "YOLOS Base. Balances speed and accuracy.",
+    "SenseTime/deformable-detr": "Deformable DETR. Improved efficiency with deformable attention.",
 }
 DEFAULT_GRADIO_PORT: int = 7860  # Default port for Gradio UI
 DEFAULT_FASTAPI_PORT: int = 8000  # Default port for FastAPI server
 PORT_RANGE: range = range(7860, 7870)  # Range of ports to try for Gradio
 MAX_PORT_ATTEMPTS: int = 10  # Maximum attempts to find an available port
-
 # Thread-safe storage for lazy-loaded models and processors
 models: Dict[str, any] = {}
 processors: Dict[str, any] = {}
 model_lock = threading.Lock()
-
 # ------------------------------
 # Image Processing
 # ------------------------------
-
 def process_image(
     image: Optional[Image.Image],
     url: Optional[str],
@@ -74,14 +77,12 @@ def process_image(
 ) -> Union[Dict, Tuple[Optional[Image.Image], Optional[pd.DataFrame], Optional[pd.DataFrame], Optional[pd.DataFrame], str]]:
     """
     Process an image for object detection or panoptic segmentation, handling Gradio and FastAPI inputs.
-
     Args:
         image: PIL Image object from file upload (optional).
         url: URL of the image to process (optional).
         model_name: Name of the model to use (must be in VALID_MODELS).
         for_json: If True, return JSON dict for API/JSON tab; else, return tuple for Gradio Home tab.
         confidence_threshold: Minimum confidence score for detection (default: 0.5).
-
     Returns:
         For JSON: Dict with base64-encoded image, detected objects, and confidence scores.
         For Gradio: Tuple of (annotated image, objects DataFrame, unique objects DataFrame, properties DataFrame, error message).
@@ -95,16 +96,17 @@ def process_image(
         if model_name not in VALID_MODELS:
             error_msg = f"Invalid model: {model_name}. Choose from: {VALID_MODELS}"
             return {"error": error_msg} if for_json else (None, None, None, None, error_msg)
-
+        # Validate confidence threshold
+        if not 0 <= confidence_threshold <= 1:
+            error_msg = "Confidence threshold must be between 0 and 1."
+            return {"error": error_msg} if for_json else (None, None, None, None, error_msg)
         # Calculate margin threshold: (1 - confidence_threshold) / 2 + confidence_threshold
         margin_threshold = (1 - confidence_threshold) / 2 + confidence_threshold
-
         # Load image from URL if provided
         if url:
             response = requests.get(url, timeout=10)
             response.raise_for_status()
             image = Image.open(BytesIO(response.content)).convert("RGB")
-
         # Load model and processor thread-safely
         with model_lock:
             if model_name not in models:
@@ -112,32 +114,35 @@ def process_image(
                 try:
                     # Select appropriate model and processor based on model name
                     if "yolos" in model_name:
-                        models[model_name] = YolosForObjectDetection.from_pretrained(model_name)
+                        models[model_name] = YolosForObjectDetection.from_pretrained(model_name).eval().to(device)
                         processors[model_name] = YolosImageProcessor.from_pretrained(model_name)
                     elif "panoptic" in model_name:
-                        models[model_name] = DetrForSegmentation.from_pretrained(model_name)
+                        models[model_name] = DetrForSegmentation.from_pretrained(model_name).eval().to(device)
                         processors[model_name] = DetrImageProcessor.from_pretrained(model_name)
+                    elif "deformable" in model_name:
+                        models[model_name] = DeformableDetrForObjectDetection.from_pretrained(model_name).eval().to(device)
+                        processors[model_name] = DeformableDetrImageProcessor.from_pretrained(model_name)
                     else:
-                        models[model_name] = DetrForObjectDetection.from_pretrained(model_name)
+                        models[model_name] = DetrForObjectDetection.from_pretrained(model_name).eval().to(device)
                         processors[model_name] = DetrImageProcessor.from_pretrained(model_name)
                 except Exception as e:
                     error_msg = f"Failed to load model: {str(e)}"
                     logger.error(error_msg)
                     return {"error": error_msg} if for_json else (None, None, None, None, error_msg)
             model, processor = models[model_name], processors[model_name]
-
         # Prepare image for model processing
         inputs = processor(images=image, return_tensors="pt")
+        # Move inputs to device if using GPU
+        if device == "cuda":
+            inputs = {k: v.to(device) for k, v in inputs.items()}
         with torch.no_grad():
             outputs = model(**inputs)
-
         # Initialize drawing context for annotations
         draw = ImageDraw.Draw(image)
         object_names: List[str] = []
         confidence_scores: List[float] = []
         object_counter = Counter()
         target_sizes = torch.tensor([image.size[::-1]])
-
         # Process results based on model type (panoptic or object detection)
         if "panoptic" in model_name:
             # Handle panoptic segmentation
@@ -161,29 +166,26 @@ def process_image(
                 if score > confidence_threshold:
                     object_names.append(label_name)
                     confidence_scores.append(float(score))
-                    object_counter[label_name] = float(score)
+                    object_counter[label_name] = max(object_counter.get(label_name, 0.0), float(score))
         else:
             # Handle object detection
-            results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]
+            results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=confidence_threshold)[0]
             for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-                if score > confidence_threshold:
-                    x, y, x2, y2 = box.tolist()
-                    label_name = model.config.id2label.get(label.item(), "Unknown")
-                    text = f"{label_name}: {score:.2f}"
-                    text_bbox = draw.textbbox((0, 0), text)
-                    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
-                    # Use yellow for confidence_threshold <= score < margin_threshold, green for >= margin_threshold
-                    color = "#FFFF00" if score < margin_threshold else "#32CD32"
-                    draw.rectangle([x, y, x2, y2], outline=color, width=2)
-                    draw.text((x2 - text_width - 2, y - text_height - 2), text, fill=color)
-                    object_names.append(label_name)
-                    confidence_scores.append(float(score))
-                    object_counter[label_name] = float(score)
-
+                x, y, x2, y2 = box.tolist()
+                label_name = model.config.id2label.get(label.item(), "Unknown")
+                text = f"{label_name}: {score:.2f}"
+                text_bbox = draw.textbbox((0, 0), text)
+                text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
+                # Use yellow for confidence_threshold <= score < margin_threshold, green for >= margin_threshold
+                color = "#FFFF00" if score < margin_threshold else "#32CD32"
+                draw.rectangle([x, y, x2, y2], outline=color, width=2)
+                draw.text((x2 - text_width - 2, y - text_height - 2), text, fill=color)
+                object_names.append(label_name)
+                confidence_scores.append(float(score))
+                object_counter[label_name] = max(object_counter.get(label_name, 0.0), float(score))
         # Compile unique objects and their highest confidence scores
         unique_objects = list(object_counter.keys())
         unique_confidences = [object_counter[obj] for obj in unique_objects]
-
         # Calculate image properties (metadata)
         properties: Dict[str, str] = {
             "Format": image.format if hasattr(image, "format") and image.format else "Unknown",
@@ -207,7 +209,6 @@ def process_image(
             properties["StdDev (R,G,B)"] = ", ".join(f"{s:.2f}" for s in stat.stddev)
         except Exception as e:
             logger.error(f"Error calculating image stats: {str(e)}")
-
         # Prepare output based on request type
         if for_json:
             # Return JSON with base64-encoded image
@@ -231,9 +232,8 @@ def process_image(
                 pd.DataFrame({"Unique Object": unique_objects, "Confidence Score": [f"{score:.2f}" for score in unique_confidences]})
                 if unique_objects else pd.DataFrame(columns=["Unique Object", "Confidence Score"])
             )
-            properties_df = pd.DataFrame([properties]) if properties else pd.DataFrame(columns=properties.keys())
+            properties_df = pd.DataFrame([properties]) if properties else pd.DataFrame(columns=list(properties.keys()))
             return image, objects_df, unique_objects_df, properties_df, ""
-
     except requests.RequestException as e:
         # Handle URL fetch errors
         error_msg = f"Error fetching image from URL: {str(e)}"
@@ -244,13 +244,10 @@ def process_image(
         error_msg = f"Error processing image: {str(e)}"
         logger.error(f"{error_msg}\n{traceback.format_exc()}")
         return {"error": error_msg} if for_json else (None, None, None, None, error_msg)
-
 # ------------------------------
 # FastAPI Setup
 # ------------------------------
-
 app = FastAPI(title="Object Detection API")
-
 @app.post("/detect")
 async def detect_objects_endpoint(
     file: Optional[UploadFile] = File(None),
@@ -260,16 +257,13 @@ async def detect_objects_endpoint(
 ) -> JSONResponse:
     """
     FastAPI endpoint to detect objects in an image from file upload or URL.
-
     Args:
         file: Uploaded image file (optional).
         image_url: URL of the image (optional).
         model_name: Model to use for detection (default: first VALID_MODELS entry).
         confidence_threshold: Confidence threshold for detection (default: 0.5).
-
     Returns:
         JSONResponse with base64-encoded image, detected objects, and confidence scores.
-
     Raises:
         HTTPException: For invalid inputs or processing errors.
     """
@@ -277,9 +271,6 @@ async def detect_objects_endpoint(
         # Validate input: ensure exactly one of file or URL
         if (file is None and not image_url) or (file is not None and image_url):
             raise HTTPException(status_code=400, detail="Provide either an image file or an image URL, not both.")
-        # Validate confidence threshold
-        if not 0 <= confidence_threshold <= 1:
-            raise HTTPException(status_code=400, detail="Confidence threshold must be between 0 and 1.")
         # Load image from file if provided
         image = None
         if file:
@@ -297,18 +288,14 @@ async def detect_objects_endpoint(
     except Exception as e:
         logger.error(f"Error in FastAPI endpoint: {str(e)}\n{traceback.format_exc()}")
         raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
-
 # ------------------------------
 # Gradio UI Setup
 # ------------------------------
-
 def create_gradio_ui() -> gr.Blocks:
     """
     Create and configure the Gradio UI for object detection with Home, JSON, and Help tabs.
-
     Returns:
         Gradio Blocks object representing the UI.
-
     Raises:
         RuntimeError: If UI creation fails.
     """
@@ -319,11 +306,11 @@ def create_gradio_ui() -> gr.Blocks:
             gr.Markdown(
                 f"""
                 # 🚀 Object Detection App
-                Upload an image or provide a URL to detect objects using transformer models (DETR, YOLOS).
+                Upload an image or provide a URL to detect objects using transformer models (DETR, YOLOS, Deformable DETR).
                 Running on port: {os.getenv('GRADIO_SERVER_PORT', 'auto-selected')}
+                Device: {device.upper()}
                 """
             )
-
             # Create tabbed interface
             with gr.Tabs():
                 # Home tab (formerly Image Upload)
@@ -335,6 +322,8 @@ def create_gradio_ui() -> gr.Blocks:
                             # Model selection dropdown
                             model_choice = gr.Dropdown(choices=VALID_MODELS, value=VALID_MODELS[0], label="🔎 Select Model")
                             model_info = gr.Markdown(f"**Model Info**: {MODEL_DESCRIPTIONS[VALID_MODELS[0]]}")
+                            # Confidence threshold slider
+                            confidence_slider = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.01, label="Confidence Threshold")
                             # Image upload input
                             image_input = gr.Image(type="pil", label="📷 Upload Image")
                             # Image URL input
@@ -343,14 +332,12 @@ def create_gradio_ui() -> gr.Blocks:
                             with gr.Row():
                                 submit_btn = gr.Button("✨ Detect", variant="primary")
                                 clear_btn = gr.Button("🗑️ Clear", variant="secondary")
-
                             # Update model info when model changes
                             model_choice.change(
                                 fn=lambda model_name: f"**Model Info**: {MODEL_DESCRIPTIONS.get(model_name, 'No description available.')}",
                                 inputs=model_choice,
                                 outputs=model_info,
                             )
-
                         # Right column for results
                         with gr.Column(scale=2):
                             gr.Markdown("### Results")
@@ -364,21 +351,18 @@ def create_gradio_ui() -> gr.Blocks:
                                 unique_objects_output = gr.DataFrame(label="🔍 Unique Objects", interactive=False)
                             # Image properties table
                             properties_output = gr.DataFrame(label="📄 Image Properties", interactive=False)
-
                     # Process image when Detect button is clicked
                     submit_btn.click(
                         fn=process_image,
-                        inputs=[image_input, image_url_input, model_choice],
+                        inputs=[image_input, image_url_input, model_choice, confidence_slider],
                         outputs=[output_image, objects_output, unique_objects_output, properties_output, error_output],
                     )
-
                     # Clear all inputs and outputs
                     clear_btn.click(
-                        fn=lambda: [None, "", None, None, None, None],
+                        fn=lambda: [None, "", VALID_MODELS[0], 0.5, None, pd.DataFrame(columns=["Object", "Confidence Score"]), pd.DataFrame(columns=["Unique Object", "Confidence Score"]), pd.DataFrame(), None],
                         inputs=None,
-                        outputs=[image_input, image_url_input, output_image, objects_output, unique_objects_output, properties_output, error_output],
+                        outputs=[image_input, image_url_input, model_choice, confidence_slider, output_image, objects_output, unique_objects_output, properties_output, error_output],
                     )
-
                 # JSON tab for API-like output
                 with gr.Tab("🔗 JSON"):
                     with gr.Row():
@@ -388,64 +372,64 @@ def create_gradio_ui() -> gr.Blocks:
                             # Model selection dropdown
                             url_model_choice = gr.Dropdown(choices=VALID_MODELS, value=VALID_MODELS[0], label="🔎 Select Model")
                             url_model_info = gr.Markdown(f"**Model Info**: {MODEL_DESCRIPTIONS[VALID_MODELS[0]]}")
+                            # Confidence threshold slider
+                            confidence_slider_json = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.01, label="Confidence Threshold")
                             # Image upload input
                             image_input_json = gr.Image(type="pil", label="📷 Upload Image")
                             # Image URL input
                             image_url_input_json = gr.Textbox(label="🔗 Image URL", placeholder="https://example.com/image.jpg")
                             # Process button
                             url_submit_btn = gr.Button("🔄 Process", variant="primary")
-
+                            url_clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                             # Update model info when model changes
                             url_model_choice.change(
                                 fn=lambda model_name: f"**Model Info**: {MODEL_DESCRIPTIONS.get(model_name, 'No description available.')}",
                                 inputs=url_model_choice,
                                 outputs=url_model_info,
                             )
-
                         # Right column for JSON output
                         with gr.Column(scale=1):
                             # JSON output display
                             url_output = gr.JSON(label="API Response")
-
                     # Process image and return JSON when Process button is clicked
                     url_submit_btn.click(
-                        fn=lambda img, url, model: process_image(img, url, model, for_json=True),
-                        inputs=[image_input_json, image_url_input_json, url_model_choice],
+                        fn=lambda img, url, model, conf: process_image(img, url, model, for_json=True, confidence_threshold=conf),
+                        inputs=[image_input_json, image_url_input_json, url_model_choice, confidence_slider_json],
                         outputs=[url_output],
                     )
-
+                    # Clear inputs and output
+                    url_clear_btn.click(
+                        fn=lambda: [None, "", VALID_MODELS[0], 0.5, {}],
+                        inputs=None,
+                        outputs=[image_input_json, image_url_input_json, url_model_choice, confidence_slider_json, url_output],
+                    )
                 # Help tab with usage instructions
                 with gr.Tab("ℹ️ Help"):
                     gr.Markdown(
                         """
                         ## How to Use
-                        - **Home**: Select a model, upload an image or provide a URL, click "Detect" to see results.
-                        - **JSON**: Select a model, upload an image or enter a URL, click "Process" for JSON output.
-                        - **Models**: Choose DETR (detection or panoptic) or YOLOS (lightweight detection).
-                        - **Clear**: Reset inputs/outputs in Home tab.
+                        - **Home**: Select a model and confidence threshold, upload an image or provide a URL, click "Detect" to see results.
+                        - **JSON**: Select a model and confidence threshold, upload an image or enter a URL, click "Process" for JSON output.
+                        - **Models**: Choose DETR (detection or panoptic), YOLOS (lightweight detection), or Deformable DETR (efficient attention).
+                        - **Clear**: Reset inputs/outputs in Home or JSON tab.
                         - **Errors**: Check error box (Home) or JSON response (JSON) for issues.
-
                         ## Tips
                         - Use high-quality images for better results.
                         - Panoptic models provide segmentation masks for complex scenes.
                         - YOLOS-Tiny is faster for resource-constrained devices.
+                        - Adjust confidence threshold to filter detections (higher = fewer but more confident).
                         """
                     )
-
         return demo
-
     except Exception as e:
         logger.error(f"Error creating Gradio UI: {str(e)}\n{traceback.format_exc()}")
         raise RuntimeError(f"Failed to create Gradio UI: {str(e)}")
-
 # ------------------------------
 # Launcher
 # ------------------------------
-
 def parse_args() -> argparse.Namespace:
     """
     Parse command-line arguments for configuring the application.
-
     Returns:
         Parsed arguments as a Namespace object.
     """
@@ -464,16 +448,13 @@ def parse_args() -> argparse.Namespace:
     if not 0 <= args.confidence_threshold <= 1:
         parser.error("Confidence threshold must be between 0 and 1.")
     return args
-
 def find_available_port(start_port: int, port_range: range, max_attempts: int) -> Optional[int]:
     """
     Find an available port within the specified range.
-
     Args:
         start_port: Initial port to try.
         port_range: Range of ports to attempt.
         max_attempts: Maximum number of ports to try.
-
     Returns:
         Available port number, or None if no port is found.
     """
@@ -497,11 +478,9 @@ def find_available_port(start_port: int, port_range: range, max_attempts: int) -
                     raise
     logger.error(f"No available port in range {min(port_range)}-{max(port_range)}")
     return None
-
 def main() -> None:
     """
     Launch the Gradio UI and optional FastAPI server.
-
     Raises:
         SystemExit: On interruption or critical errors.
     """
@@ -516,7 +495,6 @@ def main() -> None:
         if gradio_port is None:
             logger.error("Failed to find an available port for Gradio UI")
             sys.exit(1)
-
         # Start FastAPI server in a thread if enabled
         if args.enable_fastapi:
             logger.info(f"Starting FastAPI on port {args.fastapi_port}")
@@ -525,18 +503,15 @@ def main() -> None:
                 daemon=True
             )
             fastapi_thread.start()
-
         # Launch Gradio UI
         logger.info(f"Starting Gradio UI on port {gradio_port}")
         demo = create_gradio_ui()
         demo.launch(server_port=gradio_port, server_name="0.0.0.0")
-
     except KeyboardInterrupt:
         logger.info("Application terminated by user.")
         sys.exit(0)
     except Exception as e:
         logger.error(f"Error: {str(e)}\n{traceback.format_exc()}")
         sys.exit(1)
-
 if __name__ == "__main__":
     main()