Update app.py

app.py

@@ -1,3 +1,91 @@
+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"
+    
+    # Process the image
+    result_img = img_array.copy()
+    detections = {'faces': 0, 'plates': 0, 'text': 0}  # Added text tracking
+    
+    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)}"
+    
+    message = f"Detected and blurred {detections['faces']} faces, {detections['plates']} license plates, and {detections['text']} text regions"
+    return result_img, message
+
+# Use Blocks for better layout control
 with gr.Blocks(title="Privacy Protection") as demo:
     gr.Markdown("# Privacy Protection")
     gr.Markdown(