Update app.py

app.py

@@ -1,124 +1,50 @@
-import gradio as gr
-import os
-import cv2
-import numpy as np
-import torch
-from PIL import Image
-from ultralytics import YOLO
-
-# Set environment variables for temporary directories
-os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
-os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics"
-
-# Load model
-MODEL_PATH = os.path.join(os.path.dirname(__file__), "models", "unified_detector.pt")
-
-try:
-    print(f"CUDA Available: {torch.cuda.is_available()}")
-    model = YOLO(MODEL_PATH)
-    print("Model loaded successfully")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    model = None
-
-def detect_and_blur(image):
-    """Process and blur sensitive content in image with blur for better protection"""
-    # Handle different input types
-    if isinstance(image, np.ndarray):
-        img_array = image
-    elif isinstance(image, str):
-        img_array = np.array(Image.open(image).convert('RGB'))
-    else:
-        return None, "Invalid input type"
+with gr.Blocks(title="Privacy Protection") as demo:
+    gr.Markdown("# Privacy Protection")
+    gr.Markdown(
+        "Upload an image to automatically blur faces, license plates, and text for privacy protection.\n"
+        "This model uses YOLOv8 to detect and apply blur to sensitive content."
+    )
     
-    # Process the image
-    result_img = img_array.copy()
-    detections = {'faces': 0, 'plates': 0, 'text': 0}  # Added text tracking
+    # Place examples at the top
+    with gr.Row():
+        example_images = gr.Examples(
+            examples=[
+                ['examples/example0.jpg'],
+                ['examples/example1.jpg'],
+                ['examples/example2.jpg'],
+                ['examples/example3.jpg'],
+                ['examples/example4.jpg'],
+                ['examples/example5.jpg'],
+                ['examples/example6.jpg'],
+                ['examples/example7.jpg'],
+            ],
+            inputs="image_input",
+            label="Example Images",
+        )
     
-    try:
-        # Run inference with model
-        results = model.predict(img_array, conf=0.45)  # Slightly lower threshold for better recall
-        for r in results:
-            for box in r.boxes:
-                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
-                cls_id = int(box.cls[0])
-                conf = float(box.conf[0])
-                
-                # Log detection details for debugging
-                class_name = ["license_plate", "face", "text"][cls_id] if cls_id < 3 else f"unknown({cls_id})"
-                print(f"Detected {class_name} with confidence {conf:.2f} at coordinates [{x1}, {y1}, {x2}, {y2}]")
-                
-                # Ensure coordinates are within image bounds
-                x1, y1 = max(0, x1), max(0, y1)
-                x2, y2 = min(img_array.shape[1], x2), min(img_array.shape[0], y2)
-                
-                if x2 <= x1 or y2 <= y1:
-                    continue
-                
-                # Apply stronger blur based on class
-                region = result_img[y1:y2, x1:x2]
-                
-                # Increase kernel size for more extreme blur
-                # Base kernel size on region dimensions but make it larger
-                dim_min = min(x2-x1, y2-y1)
-                kernel_size = min(99, max(25, dim_min // 2 * 2 + 1))  # Ensure odd number, max 99
-                
-                # Higher sigma for more extreme blur
-                sigma = 70 if cls_id == 1 else 85  # Higher sigma for plates and text
-                
-                if kernel_size >= 3:
-                    # Apply more extreme blur with higher sigma value
-                    blurred = cv2.GaussianBlur(region, (kernel_size, kernel_size), sigma)
-                    result_img[y1:y2, x1:x2] = blurred
-                    
-                    # Update detection counters for all three classes
-                    if cls_id == 0:
-                        detections['plates'] += 1
-                    elif cls_id == 1:
-                        detections['faces'] += 1
-                    elif cls_id == 2:
-                        detections['text'] += 1
-    except Exception as e:
-        print(f"Error during detection: {str(e)}")
-        return img_array, f"Error: {str(e)}"
+    # Split the interface into two columns
+    with gr.Row():
+        # Left column for input
+        with gr.Column(scale=1):
+            image_input = gr.Image(label="Input Image")
+            process_btn = gr.Button("Process Image", variant="primary")
+        
+        # Right column for output
+        with gr.Column(scale=1):
+            # Set a fixed height for the output image
+            image_output = gr.Image(label="Processed Image", height=400)
+            output_text = gr.Textbox(label="Detection Results")
     
-    message = f"Detected and blurred {detections['faces']} faces, {detections['plates']} license plates, and {detections['text']} text regions"
-    return result_img, message
-
-# Define the Gradio interface
-title = "Privacy Protection"
-description = (
-    "Upload an image to automatically blur faces, license plates, and text for privacy protection.\n"
-    "This model uses YOLOv8 to detect and apply blur to sensitive content."
-)
-
-# Interface with flagging disabled
-interface = gr.Interface(
-    fn=detect_and_blur,
-    inputs=gr.Image(),
-    outputs=[
-        gr.Image(label="Processed Image"),
-        gr.Textbox(label="Detection Results")
-    ],
-    title=title,
-    description=description,
-    examples=[
-        ['examples/example0.jpg'],
-        ['examples/example1.jpg'],
-        ['examples/example2.jpg'],
-        ['examples/example3.jpg'],
-        ['examples/example4.jpg'],
-        ['examples/example5.jpg'],
-        ['examples/example6.jpg'],
-        ['examples/example7.jpg'],
-    ],
-    allow_flagging="never",  # Disable flagging
-    analytics_enabled=False  # Disable analytics
-)
+    # Set up the processing function
+    process_btn.click(
+        fn=detect_and_blur,
+        inputs=[image_input],
+        outputs=[image_output, output_text]
+    )
 
 # Launch the app with the correct server parameters
 if __name__ == "__main__":
-    interface.launch(
+    demo.launch(
         server_name="0.0.0.0", 
         server_port=7860,
         share=False