Update app.py

app.py

@@ -5,11 +5,22 @@ import requests
 import io
 import os
 import spaces
+import json
+import re
 
-# Initialize object detection using Hugging Face Inference API
-class ObjectDetector:
+# Initialize object detection using the most advanced YOLO model
+class AdvancedYOLODetector:
     def __init__(self):
-        self.api_url = "https://api-inference.huggingface.co/models/facebook/detr-resnet-50"
+        # Using the most advanced YOLO model available on Hugging Face
+        # YOLOv8 is the latest and most advanced version
+        self.api_url = "https://api-inference.huggingface.co/models/ultralytics/yolov8x"
+        # Fallback models in order of preference:
+        self.fallback_models = [
+            "https://api-inference.huggingface.co/models/ultralytics/yolov8l",
+            "https://api-inference.huggingface.co/models/ultralytics/yolov8m", 
+            "https://api-inference.huggingface.co/models/ultralytics/yolov8s",
+            "https://api-inference.huggingface.co/models/ultralytics/yolov8n"
+        ]
         
     def detect(self, image, hf_token=None):
         import base64
@@ -26,63 +37,218 @@ class ObjectDetector:
     
         # Convert PIL image to base64 string
         img_buffer = io.BytesIO()
-        image.save(img_buffer, format='JPEG')
+        image.save(img_buffer, format='JPEG', quality=95)
         img_bytes = img_buffer.getvalue()
         img_base64 = base64.b64encode(img_bytes).decode("utf-8")
     
-        # Send JSON payload with base64 image
-        response = requests.post(self.api_url, headers=headers, json={"inputs": img_base64})
-    
-        if response.status_code != 200:
-            raise Exception(f"Object detection API error: {response.status_code} - {response.text}")
-    
-        return response.json()
-
+        payload = {"inputs": img_base64}
+        
+        # Try main model first, then fallbacks
+        models_to_try = [self.api_url] + self.fallback_models
+        
+        for model_url in models_to_try:
+            try:
+                response = requests.post(
+                    model_url, 
+                    headers=headers, 
+                    json=payload,
+                    timeout=45
+                )
+                
+                if response.status_code == 503:
+                    # Model is loading, wait and retry once
+                    import time
+                    time.sleep(15)
+                    response = requests.post(
+                        model_url, 
+                        headers=headers, 
+                        json=payload,
+                        timeout=45
+                    )
+                
+                if response.status_code == 200:
+                    result = response.json()
+                    if isinstance(result, list) and len(result) > 0:
+                        return result
+                    elif isinstance(result, dict) and 'error' not in result:
+                        return []
+                
+                # If this model failed, try next one
+                print(f"Model {model_url} failed with status {response.status_code}, trying next...")
+                continue
+                
+            except requests.exceptions.Timeout:
+                print(f"Timeout with model {model_url}, trying next...")
+                continue
+            except requests.exceptions.RequestException as e:
+                print(f"Network error with model {model_url}: {str(e)}, trying next...")
+                continue
+        
+        # If all models failed
+        raise Exception("All YOLO models failed or are unavailable. Please try again later.")
 
-object_detector = ObjectDetector()
+object_detector = AdvancedYOLODetector()
 
-# COCO class names for object detection
-COCO_CLASSES = [
-    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
-    'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat',
-    'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
-    'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
-    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
-    'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
-    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
-    'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
-    'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
-    'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
-    'toothbrush'
+# Extended object class names including common variations and synonyms
+COMMON_OBJECTS = [
+    # People and body parts
+    'person', 'people', 'human', 'man', 'woman', 'child', 'baby', 'face', 'head',
+    
+    # Animals
+    'cat', 'dog', 'bird', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
+    'lion', 'tiger', 'monkey', 'rabbit', 'mouse', 'rat', 'pig', 'goat', 'deer', 'fox',
+    
+    # Vehicles
+    'car', 'truck', 'bus', 'motorcycle', 'bicycle', 'bike', 'airplane', 'plane', 'boat', 
+    'ship', 'train', 'van', 'taxi', 'ambulance', 'fire truck', 'police car',
+    
+    # Furniture and household items
+    'chair', 'table', 'couch', 'sofa', 'bed', 'desk', 'shelf', 'cabinet', 'drawer',
+    'tv', 'television', 'laptop', 'computer', 'monitor', 'phone', 'mobile', 'tablet',
+    
+    # Food and drinks
+    'bottle', 'cup', 'glass', 'bowl', 'plate', 'fork', 'knife', 'spoon', 'banana', 'apple',
+    'orange', 'pizza', 'sandwich', 'cake', 'donut', 'hot dog', 'hamburger', 'coffee',
+    
+    # Sports and recreation
+    'ball', 'football', 'basketball', 'tennis ball', 'baseball', 'soccer ball',
+    'skateboard', 'surfboard', 'skis', 'bicycle', 'kite', 'frisbee',
+    
+    # Clothing and accessories
+    'hat', 'cap', 'glasses', 'sunglasses', 'bag', 'backpack', 'handbag', 'purse',
+    'umbrella', 'tie', 'shoe', 'boot', 'shirt', 'jacket', 'coat',
+    
+    # Tools and objects
+    'scissors', 'hammer', 'screwdriver', 'knife', 'pen', 'pencil', 'book', 'paper',
+    'clock', 'watch', 'key', 'remote', 'controller', 'camera', 'microphone',
+    
+    # Nature and outdoor
+    'tree', 'flower', 'plant', 'grass', 'rock', 'stone', 'mountain', 'cloud', 'sun',
+    'bench', 'sign', 'pole', 'fence', 'gate', 'building', 'house', 'window', 'door'
 ]
 
-def detect_objects(image, target_class, confidence_threshold, hf_token=None):
+def fuzzy_match_object(user_input, detected_labels):
+    """
+    Advanced matching function that handles synonyms, plurals, and fuzzy matching
+    """
+    user_input = user_input.lower().strip()
+    matches = []
+    
+    # Direct matching
+    for detection in detected_labels:
+        label = detection.get('label', '').lower()
+        
+        # Exact match
+        if label == user_input:
+            matches.append(detection)
+            continue
+            
+        # Handle plurals
+        if user_input.endswith('s') and label == user_input[:-1]:
+            matches.append(detection)
+            continue
+        if label.endswith('s') and user_input == label[:-1]:
+            matches.append(detection)
+            continue
+            
+        # Substring matching
+        if user_input in label or label in user_input:
+            matches.append(detection)
+            continue
+            
+        # Handle common synonyms
+        synonyms = {
+            'person': ['human', 'people', 'man', 'woman', 'individual'],
+            'car': ['vehicle', 'automobile', 'auto'],
+            'bike': ['bicycle', 'cycle'],
+            'phone': ['mobile', 'cellphone', 'smartphone'],
+            'tv': ['television', 'telly'],
+            'couch': ['sofa', 'settee'],
+            'bag': ['purse', 'handbag', 'backpack'],
+            'glasses': ['spectacles', 'eyeglasses'],
+            'plane': ['airplane', 'aircraft'],
+            'boat': ['ship', 'vessel'],
+            'dog': ['puppy', 'canine'],
+            'cat': ['kitten', 'feline']
+        }
+        
+        # Check if user input matches any synonym
+        for main_word, synonym_list in synonyms.items():
+            if (user_input == main_word and label in synonym_list) or \
+               (user_input in synonym_list and label == main_word):
+                matches.append(detection)
+                break
+    
+    return matches
+
+def detect_objects(image, target_object, confidence_threshold, hf_token=None):
     """
-    Detect objects in the image and return bounding boxes for the target class
+    Detect any object in the image using advanced YOLO and return bounding boxes
     """
     try:
-        # Use Hugging Face Inference API for object detection
+        if not target_object or not target_object.strip():
+            raise gr.Error("Please enter an object name to detect and remove")
+            
+        # Use advanced YOLO for object detection
         results = object_detector.detect(image, hf_token)
         
-        # Filter results for target class
-
-        target_detections = []
+        if not results or not isinstance(results, list):
+            return []
+        
+        # Apply confidence threshold first
+        filtered_detections = []
         for detection in results:
-            label = detection.get('label')
-            if (
-                label is not None and
-                label.lower() == target_class.lower() and 
-                detection['score'] > confidence_threshold
-            ):
-                target_detections.append(detection)
-        
-        return target_detections
+            if isinstance(detection, dict) and detection.get('score', 0) >= confidence_threshold:
+                filtered_detections.append(detection)
+        
+        # Use fuzzy matching to find target objects
+        target_detections = fuzzy_match_object(target_object, filtered_detections)
+        
+        # Process and validate bounding boxes
+        valid_detections = []
+        image_width, image_height = image.size
+        
+        for detection in target_detections:
+            box = detection.get('box', {})
+            
+            if box and all(key in box for key in ['xmin', 'ymin', 'xmax', 'ymax']):
+                # Convert coordinates
+                xmin = box['xmin']
+                ymin = box['ymin'] 
+                xmax = box['xmax']
+                ymax = box['ymax']
+                
+                # Handle normalized coordinates (0-1 range)
+                if xmax <= 1.0 and ymax <= 1.0:
+                    xmin = int(xmin * image_width)
+                    ymin = int(ymin * image_height)
+                    xmax = int(xmax * image_width)
+                    ymax = int(ymax * image_height)
+                
+                # Ensure coordinates are within bounds and valid
+                xmin = max(0, min(int(xmin), image_width))
+                ymin = max(0, min(int(ymin), image_height))
+                xmax = max(xmin, min(int(xmax), image_width))
+                ymax = max(ymin, min(int(ymax), image_height))
+                
+                # Only add if box has valid area
+                if xmax > xmin and ymax > ymin:
+                    detection_copy = detection.copy()
+                    detection_copy['box'] = {
+                        'xmin': xmin, 'ymin': ymin, 
+                        'xmax': xmax, 'ymax': ymax
+                    }
+                    valid_detections.append(detection_copy)
+        
+        return valid_detections
+        
     except Exception as e:
+        print(f"Detection error: {str(e)}")
         raise gr.Error(f"Object detection failed: {str(e)}")
 
 def create_mask_from_detections(image, detections, mask_expansion=10):
     """
-    Create a binary mask from object detections
+    Create a binary mask from object detections with smart expansion
     """
     width, height = image.size
     mask = Image.new('L', (width, height), 0)  # Black mask
@@ -90,11 +256,17 @@ def create_mask_from_detections(image, detections, mask_expansion=10):
     
     for detection in detections:
         box = detection['box']
-        # Expand the bounding box slightly for better masking
-        x1 = max(0, box['xmin'] - mask_expansion)
-        y1 = max(0, box['ymin'] - mask_expansion)
-        x2 = min(width, box['xmax'] + mask_expansion)
-        y2 = min(height, box['ymax'] + mask_expansion)
+        
+        # Calculate expansion based on object size
+        box_width = box['xmax'] - box['xmin']
+        box_height = box['ymax'] - box['ymin']
+        adaptive_expansion = min(mask_expansion, max(5, int(min(box_width, box_height) * 0.1)))
+        
+        # Expand the bounding box
+        x1 = max(0, box['xmin'] - adaptive_expansion)
+        y1 = max(0, box['ymin'] - adaptive_expansion)
+        x2 = min(width, box['xmax'] + adaptive_expansion)
+        y2 = min(height, box['ymax'] + adaptive_expansion)
         
         # Draw white rectangle on mask (255 = area to inpaint)
         draw.rectangle([x1, y1, x2, y2], fill=255)
@@ -102,29 +274,38 @@ def create_mask_from_detections(image, detections, mask_expansion=10):
     return mask
 
 @spaces.GPU
-def remove_objects(image, object_class, confidence_threshold, mask_expansion, inpaint_prompt, hf_token):
+def remove_objects(image, object_name, confidence_threshold, mask_expansion, inpaint_prompt, hf_token):
     """
-    Main function to remove objects from image using SDXL inpainting
+    Main function to remove any specified object from image using advanced YOLO + SDXL
     """
     try:
         if image is None:
             raise gr.Error("Please upload an image")
         
+        if not object_name or not object_name.strip():
+            raise gr.Error("Please enter the name of the object you want to remove")
+        
         # Try to get token from multiple sources
         token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
         if not token:
             raise gr.Error("Please provide your Hugging Face token or set HF_TOKEN in Space secrets")
         
-        # Step 1: Detect objects
-        detections = detect_objects(image, object_class, confidence_threshold, token)
+        # Step 1: Detect objects using advanced YOLO
+        detections = detect_objects(image, object_name, confidence_threshold, token)
         
         if not detections:
-            return image, None, f"No {object_class} objects detected with confidence > {confidence_threshold}"
+            # Provide helpful suggestions
+            suggestion_msg = f"No '{object_name}' objects detected with confidence > {confidence_threshold}.\n\n"
+            suggestion_msg += "💡 Try:\n"
+            suggestion_msg += "• Lowering the confidence threshold\n"
+            suggestion_msg += "• Using different object names (e.g., 'person' instead of 'human')\n"
+            suggestion_msg += "• Checking if the object is clearly visible in the image"
+            return image, None, suggestion_msg
         
-        # Step 2: Create mask
+        # Step 2: Create mask with adaptive expansion
         mask = create_mask_from_detections(image, detections, mask_expansion)
         
-        # Step 3: Use SDXL for inpainting via Hugging Face Inference API
+        # Step 3: Use SDXL for inpainting
         inpaint_api_url = "https://api-inference.huggingface.co/models/diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
         
         headers = {"Authorization": f"Bearer {token}"}
@@ -144,29 +325,37 @@ def remove_objects(image, object_class, confidence_threshold, mask_expansion, in
             'mask': ('mask.png', mask_bytes, 'image/png')
         }
         
+        # Enhanced inpainting prompt
+        enhanced_prompt = f"{inpaint_prompt}, photorealistic, high quality, detailed, natural lighting"
+        
         data = {
-            'prompt': inpaint_prompt,
-            'negative_prompt': 'blurry, low quality, distorted, artifacts',
-            'num_inference_steps': 20,
+            'prompt': enhanced_prompt,
+            'negative_prompt': 'blurry, low quality, distorted, artifacts, unrealistic, pixelated, noise',
+            'num_inference_steps': 25,
             'guidance_scale': 7.5,
             'strength': 0.99
         }
         
         try:
-            response = requests.post(inpaint_api_url, headers=headers, files=files, data=data, timeout=60)
+            response = requests.post(inpaint_api_url, headers=headers, files=files, data=data, timeout=90)
             
             if response.status_code == 200:
                 result_image = Image.open(io.BytesIO(response.content))
-                status_msg = f"✅ Successfully removed {len(detections)} {object_class} object(s)"
+                detected_labels = [d.get('label', 'unknown') for d in detections]
+                status_msg = f"✅ Successfully removed {len(detections)} '{object_name}' object(s)\n"
+                status_msg += f"🎯 Detected as: {', '.join(detected_labels)}\n"
+                status_msg += f"🔧 Used: Advanced YOLO + SDXL Inpainting"
             else:
                 # Fallback: return original with mask overlay for debugging
                 result_image = create_mask_overlay(image, mask)
-                status_msg = f"⚠️ SDXL inpainting failed (HTTP {response.status_code}). Showing detected areas in red."
+                status_msg = f"⚠️ SDXL inpainting failed (HTTP {response.status_code}). Showing detected areas in red.\n"
+                status_msg += f"🎯 Found {len(detections)} '{object_name}' object(s) - detection was successful"
                 
         except Exception as e:
             # Fallback: return original with mask overlay for debugging
             result_image = create_mask_overlay(image, mask)
-            status_msg = f"⚠️ SDXL inpainting failed: {str(e)}. Showing detected areas in red."
+            status_msg = f"⚠️ SDXL inpainting failed: {str(e)}. Showing detected areas in red.\n"
+            status_msg += f"🎯 Found {len(detections)} '{object_name}' object(s) - detection was successful"
         
         return result_image, mask, status_msg
         
@@ -187,19 +376,20 @@ def create_mask_overlay(image, mask):
 # Create Gradio interface
 with gr.Blocks(
     fill_height=True, 
-    title="Object Removal with SDXL",
+    title="Advanced Object Removal with YOLOv8",
     theme=gr.themes.Soft()
 ) as demo:
     
     gr.Markdown("""
-    # 🎯 Object Removal using AI Detection + SDXL Inpainting
+    # 🚀 Advanced Object Removal using YOLOv8 + SDXL Inpainting
     
-    Upload an image, specify an object class to remove, and let AI intelligently remove it!
+    Upload an image and specify **ANY object** you want to remove - no limitations!
     
     **How it works:**
-    1. 🔍 **Object Detection**: Uses Facebook's DETR model to find objects
-    2. 🎭 **Mask Generation**: Creates precise removal masks  
-    3. 🎨 **AI Inpainting**: Uses SDXL to intelligently fill the removed areas
+    1. 🔍 **YOLOv8 Detection**: Uses the most advanced YOLO model for object detection
+    2. 🧠 **Smart Matching**: Handles synonyms, plurals, and fuzzy object name matching
+    3. 🎭 **Adaptive Masking**: Creates intelligent removal masks  
+    4. 🎨 **SDXL Inpainting**: Uses state-of-the-art AI to fill removed areas seamlessly
     """)
     
     with gr.Row():
@@ -213,36 +403,46 @@ with gr.Blocks(
                 height=300
             )
             
-            object_class = gr.Dropdown(
-                choices=COCO_CLASSES,
+            object_name = gr.Textbox(
                 label="🎯 Object to Remove",
+                placeholder="Enter any object name (e.g., person, car, dog, bottle, tree, sign...)",
                 value="person",
-                allow_custom_value=True,
-                info="Select or type the object class to remove"
+                info="Type ANY object name - supports synonyms and variations!"
             )
             
+            # Add suggestions
+            with gr.Row():
+                gr.Examples(
+                    examples=[
+                        ["person"], ["car"], ["dog"], ["cat"], ["bottle"], 
+                        ["chair"], ["tree"], ["sign"], ["bag"], ["phone"]
+                    ],
+                    inputs=[object_name],
+                    label="💡 Quick Examples"
+                )
+            
             with gr.Accordion("⚙️ Advanced Settings", open=False):
                 confidence_threshold = gr.Slider(
                     minimum=0.1,
                     maximum=1.0,
-                    value=0.5,
-                    step=0.1,
+                    value=0.3,
+                    step=0.05,
                     label="🎚️ Detection Confidence",
-                    info="Higher = fewer but more confident detections"
+                    info="Lower = more detections, higher = fewer but more confident"
                 )
                 
                 mask_expansion = gr.Slider(
                     minimum=0,
                     maximum=50,
-                    value=15,
+                    value=20,
                     step=5,
                     label="📏 Mask Expansion (pixels)",
-                    info="Expand mask around detected objects"
+                    info="Expand mask around detected objects for better removal"
                 )
                 
                 inpaint_prompt = gr.Textbox(
                     label="✨ Inpainting Prompt",
-                    value="natural background, seamless, high quality, photorealistic",
+                    value="natural background, seamless, realistic environment",
                     placeholder="Describe what should replace the removed object",
                     info="Be specific about the desired background/replacement"
                 )
@@ -273,9 +473,9 @@ with gr.Blocks(
                 )
             
             status_text = gr.Textbox(
-                label="📊 Status",
+                label="📊 Status & Detection Info",
                 interactive=False,
-                max_lines=3
+                max_lines=4
             )
     
     # Event handlers
@@ -283,7 +483,7 @@ with gr.Blocks(
         fn=remove_objects,
         inputs=[
             input_image,
-            object_class,
+            object_name,
             confidence_threshold,
             mask_expansion,
             inpaint_prompt,
@@ -299,40 +499,40 @@ with gr.Blocks(
             ## 📚 Instructions
             
             1. **Upload an image** containing objects you want to remove
-            2. **Select the object class** from the dropdown (e.g., 'person', 'car', 'bottle')
+            2. **Enter ANY object name** in the text box - no restrictions!
             3. **Adjust settings** if needed:
-               - **Confidence**: Lower = more detections, higher = fewer but more accurate
-               - **Mask expansion**: Larger values remove more area around objects
-               - **Inpainting prompt**: Describe the desired replacement (e.g., "grass field", "brick wall")
-            4. **Click "Remove Objects"** and wait for processing
+               - **Confidence**: Start with 0.3, increase if too many false detections
+               - **Mask expansion**: Larger values ensure complete object removal
+               - **Inpainting prompt**: Describe the desired replacement scene
+            4. **Click "Remove Objects"** and wait for AI processing
             
-            ### 💡 Tips for Better Results:
-            - Use clear, high-resolution images
-            - Be specific in inpainting prompts: "blue sky with clouds" vs "background"
-            - For complex scenes, try different confidence thresholds
-            - Objects partially cut off at image edges may not be detected well
+            ### 💡 Smart Object Recognition:
+            - **Handles variations**: "car" = "vehicle" = "automobile"
+            - **Plural support**: "person" matches "people"
+            - **Common synonyms**: "phone" = "mobile" = "smartphone"
+            - **Fuzzy matching**: Partial name matches work too!
             """)
         
         with gr.Column():
             gr.Markdown("""
-            ## 🎯 Supported Objects
-            
-            **People & Animals**: person, cat, dog, horse, bird, cow, sheep, etc.
-            
-            **Vehicles**: car, bicycle, motorcycle, bus, truck, boat, airplane
-            
-            **Furniture**: chair, couch, bed, dining table, tv, laptop
+            ## 🎯 What Can Be Removed?
             
-            **Objects**: bottle, cup, book, phone, backpack, umbrella
+            **✅ ANY Object You Can Think Of!**
             
-            **And 60+ more COCO classes!**
+            **Popular Examples:**
+            - **People**: person, human, man, woman, child, face
+            - **Animals**: dog, cat, bird, horse, any animal name
+            - **Vehicles**: car, truck, bike, plane, boat, motorcycle
+            - **Objects**: bottle, bag, phone, chair, table, sign
+            - **Nature**: tree, flower, rock, cloud, mountain
+            - **And literally thousands more!**
             
-            ### ⚠️ Important Notes:
-            - **Token Required**: Either enter your HF token above OR set `HF_TOKEN` in Space secrets
-            - **Get Token**: Visit https://huggingface.co/settings/tokens to create one
-            - Processing may take 30-60 seconds depending on image size
-            - Results depend on object detection accuracy and image complexity
-            - Red overlay indicates detected areas when inpainting fails
+            ### ⚠️ System Info:
+            - **🚀 Powered by**: YOLOv8x (most advanced YOLO model)
+            - **🎨 Inpainting**: SDXL for photorealistic results
+            - **⏱️ Processing**: 30-90 seconds depending on complexity
+            - **🔧 Fallback**: Multiple YOLO models for reliability
+            - **Token Required**: HF token needed for API access
             """)
 
 if __name__ == "__main__":