Update app.py

app.py

@@ -11,19 +11,11 @@ import os
 import sys
 import io
 
-# Set up logging
-import logging
-logging.basicConfig(level=logging.INFO, 
-                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-                    stream=sys.stderr)
-logger = logging.getLogger(__name__)
+# Load Faster R-CNN model
+frcnn_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
+frcnn_model.eval()
 
-# Load models once at startup
-logger.info("Loading Faster R-CNN model...")
-rcnn_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
-rcnn_model.eval()
-
-logger.info("Loading DETR model...")
+# Load DETR model and processor
 detr_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 detr_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 
@@ -43,295 +35,168 @@ COCO_INSTANCE_CATEGORY_NAMES = [
     'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
 ]
 
-def faster_rcnn_detection(image, threshold=0.5):
-    """Detect objects using Faster R-CNN model"""
+def recommend_model(image):
+    """Provide a basic model recommendation based on image characteristics."""
+    if image is None:
+        return "Please upload an image to get a recommendation."
+    try:
+        img_array = np.array(image)
+        height, width = img_array.shape[:2]
+        pixel_variance = np.var(img_array)
+        # Basic heuristic: DETR is better for high-resolution, complex images; Faster R-CNN for smaller, simpler ones
+        if height * width > 1000 * 1000 or pixel_variance > 1000:
+            return "DETR is recommended for high-resolution or complex images."
+        else:
+            return "Faster R-CNN is recommended for smaller or simpler images."
+    except Exception as e:
+        return f"Error in recommendation: {str(e)}"
+
+def detect_objects_frcnn(image, threshold=0.5):
+    """Run Faster R-CNN detection."""
     if image is None:
-        logger.error("No image provided to Faster R-CNN detector")
-        return create_error_image("No image provided")
-    
+        blank_img = Image.new('RGB', (400, 400), color='white')
+        plt.figure(figsize=(10, 10))
+        plt.imshow(blank_img)
+        plt.text(0.5, 0.5, "No image provided", horizontalalignment='center', verticalalignment='center',
+                 transform=plt.gca().transAxes, fontsize=20)
+        plt.axis('off')
+        output_path = "frcnn_blank_output.png"
+        plt.savefig(output_path)
+        plt.close()
+        return output_path
+
     try:
-        logger.info(f"Processing image with Faster R-CNN (threshold: {threshold})")
-        
-        # Convert threshold to float
-        threshold = float(threshold)
-        
-        # Apply transforms required by the model
+        threshold = float(threshold) if threshold is not None else 0.5
         transform = FasterRCNN_ResNet50_FPN_Weights.DEFAULT.transforms()
         image_tensor = transform(image).unsqueeze(0)
-        
-        # Run detection
+
         with torch.no_grad():
-            prediction = rcnn_model(image_tensor)[0]
-        
-        # Extract results
+            prediction = frcnn_model(image_tensor)[0]
+
         boxes = prediction['boxes'].cpu().numpy()
         labels = prediction['labels'].cpu().numpy()
         scores = prediction['scores'].cpu().numpy()
-        
-        # Create visualization
+
         image_np = np.array(image)
         plt.figure(figsize=(10, 10))
         plt.imshow(image_np)
         ax = plt.gca()
-        
-        # Draw bounding boxes
+
         for box, label, score in zip(boxes, labels, scores):
             if score >= threshold:
                 x1, y1, x2, y2 = box
-                ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1,
-                            fill=False, color='red', linewidth=2))
+                ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, color='red', linewidth=2))
                 class_name = COCO_INSTANCE_CATEGORY_NAMES[label]
-                ax.text(x1, y1, f'{class_name}: {score:.2f}',
-                        bbox=dict(facecolor='yellow', alpha=0.5),
-                        fontsize=12, color='black')
-        
-        plt.title("Faster R-CNN Detection")
+                ax.text(x1, y1, f'{class_name}: {score:.2f}', bbox=dict(facecolor='yellow', alpha=0.5), fontsize=12, color='black')
+
         plt.axis('off')
         plt.tight_layout()
-        
-        # Save result to image
-        output = io.BytesIO()
-        plt.savefig(output, format='png')
+        output_path = "frcnn_output.png"
+        plt.savefig(output_path)
         plt.close()
-        output.seek(0)
-        
-        return Image.open(output)
-    
+        return output_path
     except Exception as e:
-        logger.error(f"Error in Faster R-CNN detection: {e}")
-        import traceback
-        traceback.print_exc(file=sys.stderr)
-        return create_error_image(f"Faster R-CNN error: {str(e)}")
+        error_img = Image.new('RGB', (400, 400), color='white')
+        plt.figure(figsize=(10, 10))
+        plt.imshow(error_img)
+        plt.text(0.5, 0.5, f"Error: {str(e)}", horizontalalignment='center', verticalalignment='center',
+                 transform=plt.gca().transAxes, fontsize=12, wrap=True)
+        plt.axis('off')
+        error_path = "frcnn_error_output.png"
+        plt.savefig(error_path)
+        plt.close()
+        return error_path
 
-def detr_detection(image, threshold=0.5):
-    """Detect objects using DETR model"""
+def detect_objects_detr(image, threshold=0.9):
+    """Run DETR detection."""
     if image is None:
-        logger.error("No image provided to DETR detector")
-        return create_error_image("No image provided")
-    
+        blank_img = Image.new('RGB', (400, 400), color='white')
+        fig, ax = plt.subplots(1, figsize=(10, 10))
+        ax.imshow(blank_img)
+        ax.text(0.5, 0.5, "No image provided", horizontalalignment='center', verticalalignment='center',
+                transform=ax.transAxes, fontsize=20)
+        plt.axis('off')
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+        plt.close(fig)
+        buf.seek(0)
+        return Image.open(buf)
+
     try:
-        logger.info(f"Processing image with DETR (threshold: {threshold})")
-        
-        # Convert threshold to float
-        threshold = float(threshold)
-        
-        # Process image and run model
         inputs = detr_processor(images=image, return_tensors="pt")
         outputs = detr_model(**inputs)
-        
-        # Post-process results
         target_sizes = torch.tensor([image.size[::-1]])
-        results = detr_processor.post_process_object_detection(
-            outputs, target_sizes=target_sizes, threshold=threshold)[0]
-        
-        # Create visualization
+        results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=threshold)[0]
+
         fig, ax = plt.subplots(1, figsize=(10, 10))
         ax.imshow(image)
-        
-        # Draw bounding boxes
+
         for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
             xmin, ymin, xmax, ymax = box.tolist()
-            ax.add_patch(patches.Rectangle(
-                (xmin, ymin), xmax - xmin, ymax - ymin,
-                linewidth=2, edgecolor='blue', facecolor='none'))
-            ax.text(xmin, ymin, f"{detr_model.config.id2label[label.item()]}: {score:.2f}",
-                   bbox=dict(facecolor='cyan', alpha=0.5), fontsize=12)
-        
-        plt.title("DETR Detection")
+            ax.add_patch(patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, linewidth=2, edgecolor='red', facecolor='none'))
+            ax.text(xmin, ymin, f"{detr_model.config.id2label[label.item()]}: {round(score.item(), 2)}",
+                    bbox=dict(facecolor='yellow', alpha=0.5), fontsize=8)
+
         plt.axis('off')
-        plt.tight_layout()
-        
-        # Save result to image
-        output = io.BytesIO()
-        plt.savefig(output, format='png')
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
         plt.close(fig)
-        output.seek(0)
-        
-        return Image.open(output)
-    
+        buf.seek(0)
+        return Image.open(buf)
     except Exception as e:
-        logger.error(f"Error in DETR detection: {e}")
-        import traceback
-        traceback.print_exc(file=sys.stderr)
-        return create_error_image(f"DETR error: {str(e)}")
-
-def compare_detections(image, threshold=0.5):
-    """Run both models and return side-by-side comparison"""
-    if image is None:
-        logger.error("No image provided for comparison")
-        return create_error_image("No image provided")
-    
-    try:
-        logger.info(f"Comparing both models with threshold: {threshold}")
-        
-        # Run both models
-        rcnn_result = faster_rcnn_detection(image, threshold)
-        detr_result = detr_detection(image, threshold)
-        
-        # Create side-by-side comparison
-        fig, axes = plt.subplots(1, 2, figsize=(20, 10))
-        
-        axes[0].imshow(np.array(rcnn_result))
-        axes[0].set_title("Faster R-CNN Detection", fontsize=16)
-        axes[0].axis('off')
-        
-        axes[1].imshow(np.array(detr_result))
-        axes[1].set_title("DETR Detection", fontsize=16)
-        axes[1].axis('off')
-        
-        plt.tight_layout()
-        
-        # Save comparison to image
-        output = io.BytesIO()
-        plt.savefig(output, format='png', dpi=120)
+        error_img = Image.new('RGB', (400, 400), color='white')
+        fig, ax = plt.subplots(1, figsize=(10, 10))
+        ax.imshow(error_img)
+        ax.text(0.5, 0.5, f"Error: {str(e)}", horizontalalignment='center', verticalalignment='center',
+                transform=ax.transAxes, fontsize=12, wrap=True)
+        plt.axis('off')
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
         plt.close(fig)
-        output.seek(0)
-        
-        return Image.open(output)
-    
-    except Exception as e:
-        logger.error(f"Error in comparison: {e}")
-        import traceback
-        traceback.print_exc(file=sys.stderr)
-        return create_error_image(f"Comparison error: {str(e)}")
+        buf.seek(0)
+        return Image.open(buf)
 
-def create_error_image(error_text):
-    """Create an image with error message"""
-    error_img = Image.new('RGB', (800, 400), color='white')
-    fig, ax = plt.subplots(figsize=(8, 4))
-    ax.imshow(error_img)
-    ax.text(0.5, 0.5, f"Error: {error_text}",
-            horizontalalignment='center', verticalalignment='center',
-            transform=ax.transAxes, fontsize=14, wrap=True)
-    ax.axis('off')
-    
-    # Save to buffer
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    plt.close(fig)
-    buf.seek(0)
-    
-    return Image.open(buf)
+def run_detection(image, model_choice, frcnn_threshold=0.5, detr_threshold=0.9):
+    """Run detection based on model choice and return results with recommendation."""
+    recommendation = recommend_model(image)
+    frcnn_result = None
+    detr_result = None
 
-def detect_objects(image, model_choice, threshold=0.5):
-    """Main detection function that routes to the appropriate model"""
-    if image is None:
-        return create_error_image("No image provided")
-    
-    if model_choice == "Faster R-CNN":
-        return faster_rcnn_detection(image, threshold)
-    elif model_choice == "DETR":
-        return detr_detection(image, threshold)
-    elif model_choice == "Compare Both":
-        return compare_detections(image, threshold)
-    else:
-        return create_error_image(f"Unknown model choice: {model_choice}")
+    if model_choice in ["Faster R-CNN", "Both"]:
+        frcnn_result = detect_objects_frcnn(image, frcnn_threshold)
+    if model_choice in ["DETR", "Both"]:
+        detr_result = detect_objects_detr(image, detr_threshold)
 
-def model_info(model_choice):
-    """Provide information about the selected model"""
-    if model_choice == "Faster R-CNN":
-        return """
-        **Faster R-CNN** is a two-stage object detection model that first proposes regions of interest and then classifies them.
-        
-        **Strengths:**
-        - Generally high accuracy
-        - Good for detecting objects of various sizes
-        - Well-established architecture with many pretrained variants
-        
-        **Suited for:**
-        - General object detection tasks
-        - Scenes with multiple objects of different scales
-        - When detection accuracy is more important than speed
-        """
-    elif model_choice == "DETR":
-        return """
-        **DETR (DEtection TRansformer)** is an end-to-end object detection model using transformers.
-        
-        **Strengths:**
-        - Clean, end-to-end architecture without manual anchors or NMS
-        - Strong spatial reasoning via self-attention
-        - Good at dealing with occlusion
-        
-        **Suited for:**
-        - Scenes with overlapping objects
-        - When you need global context understanding
-        - Modern transformer-based approach to detection
-        """
-    elif model_choice == "Compare Both":
-        return """
-        **Comparison Mode** runs both Faster R-CNN and DETR side by side to compare their detection results.
-        
-        This is useful for:
-        - Understanding the strengths of each model
-        - Seeing how detection approaches differ
-        - Choosing the right model for your specific use case
-        """
-    return ""
+    return recommendation, frcnn_result, detr_result
+
+# Example image paths
+examples = [
+    os.path.join("/home/user/app", "TEST_IMG_1.jpg"),
+    os.path.join("/home/user/app", "TEST_IMG_2.JPG"),
+    os.path.join("/home/user/app", "TEST_IMG_3.jpg"),
+    os.path.join("/home/user/app", "TEST_IMG_4.jpg")
+]
+example_list = [[path] for path in examples if os.path.exists(path)]
 
-# Create Gradio interface
-with gr.Blocks(title="Object Detection Model Comparison") as demo:
-    gr.Markdown("""
-    # Object Detection Model Comparison
-    
-    Upload an image and choose between two state-of-the-art object detection models:
-    - **Faster R-CNN**: A classic two-stage detector
-    - **DETR**: A modern transformer-based detector
-    
-    Adjust the confidence threshold to control detection sensitivity.
-    """)
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            # Input controls
-            input_image = gr.Image(type="pil", label="Input Image")
-            model_dropdown = gr.Dropdown(
-                choices=["Faster R-CNN", "DETR", "Compare Both"],
-                value="Compare Both",
-                label="Detection Model"
-            )
-            threshold_slider = gr.Slider(
-                minimum=0.0, maximum=1.0, value=0.5, step=0.05,
-                label="Confidence Threshold"
-            )
-            detect_button = gr.Button("Detect Objects", variant="primary")
-            
-            # Model info box
-            model_info_box = gr.Markdown()
-            
-        with gr.Column(scale=2):
-            # Output image
-            output_image = gr.Image(label="Detection Results")
-            
-    # Connect components
-    detect_button.click(
-        detect_objects,
-        inputs=[input_image, model_dropdown, threshold_slider],
-        outputs=output_image
-    )
-    
-    model_dropdown.change(
-        model_info,
-        inputs=model_dropdown,
-        outputs=model_info_box
-    )
-    
-    # Add examples
-    examples_dir = "/home/user/app"
-    examples = [
-        [os.path.join(examples_dir, "TEST_IMG_1.jpg"), "Compare Both", 0.5],
-        [os.path.join(examples_dir, "TEST_IMG_2.JPG"), "Compare Both", 0.5],
-        [os.path.join(examples_dir, "TEST_IMG_3.jpg"), "Compare Both", 0.5],
-        [os.path.join(examples_dir, "TEST_IMG_4.jpg"), "Compare Both", 0.5]
-    ]
-    
-    gr.Examples(
-        examples=examples,
-        inputs=[input_image, model_dropdown, threshold_slider],
-        outputs=output_image,
-        fn=detect_objects,
-        cache_examples=False
-    )
+# Gradio interface
+interface = gr.Interface(
+    fn=run_detection,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Dropdown(choices=["Faster R-CNN", "DETR", "Both"], label="Model Choice", value="Both"),
+        gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="Faster R-CNN Confidence Threshold"),
+        gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.05, label="DETR Confidence Threshold")
+    ],
+    outputs=[
+        gr.Textbox(label="Model Recommendation"),
+        gr.Image(type="filepath", label="Faster R-CNN Result"),
+        gr.Image(type="pil", label="DETR Result")
+    ],
+    title="Object Detection: Faster R-CNN vs DETR",
+    description="Upload an image, select a model (or both), and view object detection results. A recommendation is provided based on image characteristics.",
+    examples=example_list,
+    cache_examples=False
+)
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    interface.launch(debug=True)