Update app.py

app.py

@@ -1,22 +1,20 @@
 import gradio as gr
 import torch
-from transformers import AutoProcessor, AutoModel # Keep CLIP for potential future use or if FastSAM's text prompt isn't enough
+from transformers import AutoProcessor, AutoModel
 from PIL import Image, ImageDraw, ImageFont
 import numpy as np
 import random
 import os
-import wget # To download weights
-import traceback # For detailed error printing
+import wget
+import traceback
 
 # --- Configuration & Model Loading ---
 
-# Device Selection
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# Force CPU if CUDA fails or isn't desired (sometimes needed on Spaces free tier)
-# DEVICE = "cpu"
+# Device Selection with fallback
+DEVICE = "cuda" if torch.cuda.is_available() and torch.cuda.current_device() >= 0 else "cpu"
 print(f"Using device: {DEVICE}")
 
-# --- CLIP Setup (Kept in case needed, but FastSAM's method is primary now) ---
+# --- CLIP Setup ---
 CLIP_MODEL_ID = "openai/clip-vit-base-patch32"
 clip_processor = None
 clip_model = None
@@ -30,7 +28,7 @@ def load_clip_model():
             print("CLIP processor loaded.")
         except Exception as e:
             print(f"Error loading CLIP processor: {e}")
-            return False # Indicate failure
+            return False
     if clip_model is None:
         try:
             print(f"Loading CLIP model: {CLIP_MODEL_ID}...")
@@ -38,255 +36,218 @@ def load_clip_model():
             print(f"CLIP model loaded to {DEVICE}.")
         except Exception as e:
             print(f"Error loading CLIP model: {e}")
-            return False # Indicate failure
-    return True # Indicate success
-
+            return False
+    return True
 
 # --- FastSAM Setup ---
 FASTSAM_CHECKPOINT = "FastSAM-s.pt"
 FASTSAM_CHECKPOINT_URL = f"https://huggingface.co/CASIA-IVA-Lab/FastSAM-s/resolve/main/{FASTSAM_CHECKPOINT}"
 
 fastsam_model = None
-fastsam_lib_imported = False # Flag to check if import worked
+fastsam_lib_imported = False
 
 def check_and_import_fastsam():
     global fastsam_lib_imported
     if not fastsam_lib_imported:
         try:
             from fastsam import FastSAM, FastSAMPrompt
-            globals()['FastSAM'] = FastSAM # Make classes available globally
+            globals()['FastSAM'] = FastSAM
             globals()['FastSAMPrompt'] = FastSAMPrompt
             fastsam_lib_imported = True
             print("fastsam library imported successfully.")
-        except ImportError:
-            print("Error: 'fastsam' library not found or import failed.")
-            print("Please ensure 'fastsam' is installed correctly (pip install fastsam).")
+        except ImportError as e:
+            print(f"Error: 'fastsam' library not found. Install with 'pip install fastsam': {e}")
             fastsam_lib_imported = False
         except Exception as e:
-            print(f"An unexpected error occurred during fastsam import: {e}")
+            print(f"Unexpected error during fastsam import: {e}")
             traceback.print_exc()
             fastsam_lib_imported = False
     return fastsam_lib_imported
 
-
-def download_fastsam_weights():
+def download_fastsam_weights(retries=3):
     if not os.path.exists(FASTSAM_CHECKPOINT):
         print(f"Downloading FastSAM weights: {FASTSAM_CHECKPOINT} from {FASTSAM_CHECKPOINT_URL}...")
-        try:
-            wget.download(FASTSAM_CHECKPOINT_URL, FASTSAM_CHECKPOINT)
-            print("FastSAM weights downloaded.")
-        except Exception as e:
-            print(f"Error downloading FastSAM weights: {e}")
-            print("Please ensure the URL is correct and reachable, or manually place the weights file.")
-            if os.path.exists(FASTSAM_CHECKPOINT):
-                 try: os.remove(FASTSAM_CHECKPOINT)
-                 except OSError: pass
-            return False
+        for attempt in range(retries):
+            try:
+                wget.download(FASTSAM_CHECKPOINT_URL, FASTSAM_CHECKPOINT)
+                print("FastSAM weights downloaded.")
+                break
+            except Exception as e:
+                print(f"Attempt {attempt + 1}/{retries} failed: {e}")
+                if attempt + 1 == retries:
+                    print("Failed to download weights after all attempts.")
+                    return False
     return os.path.exists(FASTSAM_CHECKPOINT)
 
 def load_fastsam_model():
     global fastsam_model
     if fastsam_model is None:
         if not check_and_import_fastsam():
-             print("Cannot load FastSAM model because the library couldn't be imported.")
-             return False # Indicate failure
-
+            print("Cannot load FastSAM model due to library import failure.")
+            return False
         if download_fastsam_weights():
             try:
                 print(f"Loading FastSAM model: {FASTSAM_CHECKPOINT}...")
                 fastsam_model = FastSAM(FASTSAM_CHECKPOINT)
-                # The FastSAM model itself doesn't need explicit .to(DEVICE)
-                # It seems to handle device selection internally or via the prompt process
-                print(f"FastSAM model loaded.")
-                return True # Indicate success
+                print("FastSAM model loaded.")
+                return True
             except Exception as e:
                 print(f"Error loading FastSAM model: {e}")
                 traceback.print_exc()
+                return False
         else:
-            print("FastSAM weights not found or download failed. Cannot load model.")
-    return fastsam_model is not None # Return True if already loaded or loaded successfully
-
+            print("FastSAM weights not found or download failed.")
+            return False
+    return True
 
 # --- Processing Functions ---
 
-# (Keep run_clip_zero_shot and run_fastsam_segmentation as they were for the other tabs)
-# CLIP Zero-Shot Classification Function
 def run_clip_zero_shot(image: Image.Image, text_labels: str):
-    # Load CLIP if needed
     if clip_model is None or clip_processor is None:
         if not load_clip_model():
-            return "Error: CLIP Model could not be loaded. Check logs.", None
-
-    if image is None: return "Please upload an image.", None
-    if not text_labels: return {}, image # Return empty dict, show image
+            return "Error: CLIP Model could not be loaded.", None
+    if image is None:
+        return "Please upload an image.", None
+    if not text_labels:
+        return {}, image
 
     labels = [label.strip() for label in text_labels.split(',') if label.strip()]
-    if not labels: return {}, image
+    if not labels:
+        return {}, image
 
     print(f"Running CLIP zero-shot classification with labels: {labels}")
     try:
-        if image.mode != "RGB": image = image.convert("RGB")
+        if image.mode != "RGB":
+            image = image.convert("RGB")
         inputs = clip_processor(text=labels, images=image, return_tensors="pt", padding=True).to(DEVICE)
         with torch.no_grad():
             outputs = clip_model(**inputs)
             probs = outputs.logits_per_image.softmax(dim=1)
-        print("CLIP processing complete.")
         confidences = {labels[i]: float(probs[0, i].item()) for i in range(len(labels))}
         return confidences, image
     except Exception as e:
         print(f"Error during CLIP processing: {e}")
         traceback.print_exc()
-        return f"An error occurred during CLIP: {e}", image
+        return f"Error: {e}", image
 
-# FastSAM Everything Segmentation Function (for the second tab)
 def run_fastsam_segmentation(image_pil: Image.Image, conf_threshold: float = 0.4, iou_threshold: float = 0.9):
-    if not load_fastsam_model():
-        return "Error: FastSAM Model not loaded. Check logs."
-    if not fastsam_lib_imported:
-        return "Error: FastSAM library not available."
-    if image_pil is None: return "Please upload an image."
+    if not load_fastsam_model() or not fastsam_lib_imported:
+        return "Error: FastSAM not loaded or library unavailable."
+    if image_pil is None:
+        return "Please upload an image."
 
     print("Running FastSAM 'segment everything'...")
     try:
-        if image_pil.mode != "RGB": image_pil = image_pil.convert("RGB")
+        if image_pil.mode != "RGB":
+            image_pil = image_pil.convert("RGB")
         image_np_rgb = np.array(image_pil)
 
         everything_results = fastsam_model(
             image_np_rgb, device=DEVICE, retina_masks=True, imgsz=640,
-            conf=conf_threshold, iou=iou_threshold,
+            conf=conf_threshold, iou=iou_threshold, verbose=True
         )
         prompt_process = FastSAMPrompt(image_np_rgb, everything_results, device=DEVICE)
         ann = prompt_process.everything_prompt()
-        print(f"FastSAM 'everything' found {len(ann[0]['masks']) if ann and ann[0] and 'masks' in ann[0] else 0} masks.")
 
-        # Plotting
         output_image = image_pil.copy()
-        if ann and ann[0] is not None and 'masks' in ann[0] and len(ann[0]['masks']) > 0:
+        if ann and ann[0] and 'masks' in ann[0] and len(ann[0]['masks']) > 0:
             masks = ann[0]['masks'].cpu().numpy()
+            print(f"Found {len(masks)} masks with shape: {masks.shape}")
             overlay = Image.new('RGBA', output_image.size, (0, 0, 0, 0))
             draw = ImageDraw.Draw(overlay)
             for mask in masks:
-                color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 128)
-                mask_image = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
+                mask = (mask > 0).astype(np.uint8) * 255
+                color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 180)
+                mask_image = Image.fromarray(mask, mode='L')
                 draw.bitmap((0, 0), mask_image, fill=color)
             output_image = Image.alpha_composite(output_image.convert('RGBA'), overlay).convert('RGB')
-
-        print("FastSAM 'everything' processing complete.")
+        else:
+            print("No masks detected in 'segment everything' mode.")
         return output_image
-
     except Exception as e:
         print(f"Error during FastSAM 'everything' processing: {e}")
         traceback.print_exc()
-        return f"An error occurred during FastSAM 'everything': {e}"
-
+        return f"Error: {e}"
 
-# --- NEW: Text-Prompted Segmentation Function ---
 def run_text_prompted_segmentation(image_pil: Image.Image, text_prompts: str, conf_threshold: float = 0.4, iou_threshold: float = 0.9):
-    """Segments objects based on text prompts."""
     if not load_fastsam_model():
-        return "Error: FastSAM Model not loaded. Check logs.", "No prompts provided."
+        return "Error: FastSAM Model not loaded.", "Model load failure."
     if not fastsam_lib_imported:
-        return "Error: FastSAM library not available.", "FastSAM library error."
+        return "Error: FastSAM library not available.", "Library import error."
     if image_pil is None:
         return "Please upload an image.", "No image provided."
     if not text_prompts:
-        return image_pil, "Please enter text prompts (e.g., 'person, dog')." # Return original image and message
+        return image_pil, "Please enter text prompts (e.g., 'person, dog')."
 
     prompts = [p.strip() for p in text_prompts.split(',') if p.strip()]
     if not prompts:
         return image_pil, "No valid text prompts entered."
 
     print(f"Running FastSAM text-prompted segmentation for: {prompts}")
-
     try:
         if image_pil.mode != "RGB":
             image_pil = image_pil.convert("RGB")
         image_np_rgb = np.array(image_pil)
 
-        # 1. Run FastSAM once to get all potential results
-        # NOTE: We might optimize later, but this is the standard way FastSAMPrompt works.
         everything_results = fastsam_model(
             image_np_rgb, device=DEVICE, retina_masks=True, imgsz=640,
-            conf=conf_threshold, iou=iou_threshold, verbose=False # Less console spam
+            conf=conf_threshold, iou=iou_threshold, verbose=True
         )
-
-        # 2. Create the prompt processor
         prompt_process = FastSAMPrompt(image_np_rgb, everything_results, device=DEVICE)
-
-        # 3. Use text_prompt for each prompt and collect masks
         all_matching_masks = []
         found_prompts = []
 
         for text in prompts:
             print(f"  Processing prompt: '{text}'")
-            # Ann is a list of dictionaries, one per image. We have one image.
-            # Each dict can have 'masks', 'bboxes', 'points'.
-            # text_prompt filters 'everything_results' based on CLIP-like similarity.
-            # It might return multiple masks if multiple instances match the text.
             ann = prompt_process.text_prompt(text=text)
-
-            if ann and ann[0] is not None and 'masks' in ann[0] and len(ann[0]['masks']) > 0:
+            if ann and ann[0] and 'masks' in ann[0] and len(ann[0]['masks']) > 0:
                 num_found = len(ann[0]['masks'])
-                print(f"    Found {num_found} mask(s) matching '{text}'.")
+                print(f"    Found {num_found} mask(s) with shape: {ann[0]['masks'].shape}")
                 found_prompts.append(f"{text} ({num_found})")
-                masks = ann[0]['masks'].cpu().numpy() # Get masks as numpy array (N, H, W)
-                all_matching_masks.extend(masks) # Add the numpy arrays to the list
+                masks = ann[0]['masks'].cpu().numpy()
+                all_matching_masks.extend(masks)
             else:
-                print(f"    No masks found matching '{text}'.")
+                print(f"    No masks found for '{text}'.")
                 found_prompts.append(f"{text} (0)")
 
-        # 4. Plot the collected masks
         output_image = image_pil.copy()
-        status_message = f"Found segments for: {', '.join(found_prompts)}" if found_prompts else "No matching segments found for any prompt."
-
-        if not all_matching_masks:
-            print("No matching masks found for any prompt.")
-            return output_image, status_message # Return original image if nothing matched
-
-        # Convert list of (H, W) masks to a single (N, H, W) array for consistent processing
-        masks_np = np.stack(all_matching_masks, axis=0) # Shape (TotalMasks, H, W)
-
-        overlay = Image.new('RGBA', output_image.size, (0, 0, 0, 0))
-        draw = ImageDraw.Draw(overlay)
+        status_message = f"Found segments for: {', '.join(found_prompts)}" if found_prompts else "No matches found."
 
-        for i in range(masks_np.shape[0]):
-            mask = masks_np[i] # Shape (H, W), boolean
-            color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 150) # RGBA with slightly more alpha
-            mask_image = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
-            draw.bitmap((0, 0), mask_image, fill=color)
-
-        output_image = Image.alpha_composite(output_image.convert('RGBA'), overlay).convert('RGB')
+        if all_matching_masks:
+            masks_np = np.stack(all_matching_masks, axis=0)
+            print(f"Total masks stacked: {masks_np.shape}")
+            overlay = Image.new('RGBA', output_image.size, (0, 0, 0, 0))
+            draw = ImageDraw.Draw(overlay)
+            for mask in masks_np:
+                mask = (mask > 0).astype(np.uint8) * 255
+                color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255), 180)
+                mask_image = Image.fromarray(mask, mode='L')
+                draw.bitmap((0, 0), mask_image, fill=color)
+            output_image = Image.alpha_composite(output_image.convert('RGBA'), overlay).convert('RGB')
 
-        print("FastSAM text-prompted processing complete.")
         return output_image, status_message
-
     except Exception as e:
         print(f"Error during FastSAM text-prompted processing: {e}")
         traceback.print_exc()
-        return f"An error occurred: {e}", "Error during processing."
-
+        return image_pil, f"Error: {e}"
 
 # --- Gradio Interface ---
 
 print("Attempting to preload models...")
-# load_clip_model() # Load CLIP lazily if needed
-load_fastsam_model() # Load FastSAM eagerly
-print("Preloading finished (or attempted).")
-
+load_fastsam_model()  # Load FastSAM eagerly
+print("Preloading finished.")
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# CLIP & FastSAM Demo")
     gr.Markdown("Explore Zero-Shot Classification, 'Segment Everything', and Text-Prompted Segmentation.")
 
     with gr.Tabs():
-        # --- CLIP Tab (No changes) ---
         with gr.TabItem("CLIP Zero-Shot Classification"):
-            # ... (keep the existing layout and logic for CLIP) ...
-            gr.Markdown("Upload an image and provide comma-separated candidate labels (e.g., 'cat, dog, car'). CLIP will predict the probability of the image matching each label.")
+            gr.Markdown("Upload an image and provide comma-separated labels (e.g., 'cat, dog, car').")
             with gr.Row():
                 with gr.Column(scale=1):
                     clip_input_image = gr.Image(type="pil", label="Input Image")
-                    clip_text_labels = gr.Textbox(label="Comma-Separated Labels", placeholder="e.g., astronaut, moon, dog playing fetch")
+                    clip_text_labels = gr.Textbox(label="Comma-Separated Labels", placeholder="e.g., astronaut, moon")
                     clip_button = gr.Button("Run CLIP Classification", variant="primary")
                 with gr.Column(scale=1):
                     clip_output_label = gr.Label(label="Classification Probabilities")
@@ -298,70 +259,69 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             )
             gr.Examples(
                 examples=[
-                    ["examples/astronaut.jpg", "astronaut, moon, rover, mountain"],
-                    ["examples/dog_bike.jpg", "dog, bicycle, person, park, grass"],
-                    ["examples/clip_logo.png", "logo, text, graphics, abstract art"],
+                    ["examples/astronaut.jpg", "astronaut, moon, rover"],
+                    ["examples/dog_bike.jpg", "dog, bicycle, person"],
+                    ["examples/clip_logo.png", "logo, text, graphics"],
                 ],
                 inputs=[clip_input_image, clip_text_labels],
-                outputs=[clip_output_label, clip_output_image_display], fn=run_clip_zero_shot, cache_examples=False,
+                outputs=[clip_output_label, clip_output_image_display],
+                fn=run_clip_zero_shot,
+                cache_examples=False,
             )
 
-
-        # --- FastSAM Everything Tab (No changes) ---
         with gr.TabItem("FastSAM Segment Everything"):
-            # ... (keep the existing layout and logic for segment everything) ...
-             gr.Markdown("Upload an image. FastSAM will attempt to segment all objects/regions in the image.")
-             with gr.Row():
-                 with gr.Column(scale=1):
-                     fastsam_input_image_all = gr.Image(type="pil", label="Input Image", elem_id="fastsam_input_all") # Unique elem_id if needed
-                     with gr.Row():
-                         fastsam_conf_all = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold")
-                         fastsam_iou_all = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="IoU Threshold")
-                     fastsam_button_all = gr.Button("Run FastSAM Segmentation", variant="primary")
-                 with gr.Column(scale=1):
-                     fastsam_output_image_all = gr.Image(type="pil", label="Segmented Image", elem_id="fastsam_output_all")
-             fastsam_button_all.click(
-                 run_fastsam_segmentation,
-                 inputs=[fastsam_input_image_all, fastsam_conf_all, fastsam_iou_all],
-                 outputs=[fastsam_output_image_all]
-             )
-             gr.Examples(
-                 examples=[
-                     ["examples/dogs.jpg", 0.4, 0.9],
-                     ["examples/fruits.jpg", 0.5, 0.8],
-                     ["examples/lion.jpg", 0.45, 0.9],
-                 ],
-                 inputs=[fastsam_input_image_all, fastsam_conf_all, fastsam_iou_all],
-                 outputs=[fastsam_output_image_all], fn=run_fastsam_segmentation, cache_examples=False,
-             )
-
-        # --- NEW: Text-Prompted Segmentation Tab ---
+            gr.Markdown("Upload an image to segment all objects/regions.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    fastsam_input_image_all = gr.Image(type="pil", label="Input Image")
+                    with gr.Row():
+                        fastsam_conf_all = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold")
+                        fastsam_iou_all = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="IoU Threshold")
+                    fastsam_button_all = gr.Button("Run FastSAM Segmentation", variant="primary")
+                with gr.Column(scale=1):
+                    fastsam_output_image_all = gr.Image(type="pil", label="Segmented Image")
+            fastsam_button_all.click(
+                run_fastsam_segmentation,
+                inputs=[fastsam_input_image_all, fastsam_conf_all, fastsam_iou_all],
+                outputs=[fastsam_output_image_all]
+            )
+            gr.Examples(
+                examples=[
+                    ["examples/dogs.jpg", 0.4, 0.9],
+                    ["examples/fruits.jpg", 0.5, 0.8],
+                    ["examples/lion.jpg", 0.45, 0.9],
+                ],
+                inputs=[fastsam_input_image_all, fastsam_conf_all, fastsam_iou_all],
+                outputs=[fastsam_output_image_all],
+                fn=run_fastsam_segmentation,
+                cache_examples=False,
+            )
+
         with gr.TabItem("Text-Prompted Segmentation"):
-            gr.Markdown("Upload an image and provide comma-separated text prompts (e.g., 'person, dog, backpack'). FastSAM + CLIP (internally) will segment only the objects matching the text.")
+            gr.Markdown("Upload an image and provide comma-separated prompts (e.g., 'person, dog').")
             with gr.Row():
                 with gr.Column(scale=1):
                     prompt_input_image = gr.Image(type="pil", label="Input Image")
-                    prompt_text_input = gr.Textbox(label="Comma-Separated Text Prompts", placeholder="e.g., glasses, watch, t-shirt")
-                    with gr.Row(): # Reuse confidence/IoU sliders if desired
+                    prompt_text_input = gr.Textbox(label="Comma-Separated Text Prompts", placeholder="e.g., glasses, watch")
+                    with gr.Row():
                         prompt_conf = gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Confidence Threshold")
                         prompt_iou = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="IoU Threshold")
                     prompt_button = gr.Button("Segment by Text", variant="primary")
                 with gr.Column(scale=1):
                     prompt_output_image = gr.Image(type="pil", label="Text-Prompted Segmentation")
-                    prompt_status_message = gr.Textbox(label="Status", interactive=False) # To show which prompts matched
-
+                    prompt_status_message = gr.Textbox(label="Status", interactive=False)
             prompt_button.click(
                 run_text_prompted_segmentation,
                 inputs=[prompt_input_image, prompt_text_input, prompt_conf, prompt_iou],
-                outputs=[prompt_output_image, prompt_status_message] # Map to image and status box
+                outputs=[prompt_output_image, prompt_status_message]
             )
             gr.Examples(
                 examples=[
                     ["examples/dog_bike.jpg", "person, bicycle", 0.4, 0.9],
                     ["examples/astronaut.jpg", "person, helmet", 0.35, 0.9],
-                    ["examples/dogs.jpg", "dog", 0.4, 0.9], # Should find multiple dogs
+                    ["examples/dogs.jpg", "dog", 0.4, 0.9],
                     ["examples/fruits.jpg", "banana, apple", 0.5, 0.8],
-                    ["examples/teacher.jpg", "person, glasses, blackboard", 0.4, 0.9], # Download this image or use another one with glasses/blackboard
+                    ["examples/teacher.jpg", "person, glasses", 0.4, 0.9],
                 ],
                 inputs=[prompt_input_image, prompt_text_input, prompt_conf, prompt_iou],
                 outputs=[prompt_output_image, prompt_status_message],
@@ -369,31 +329,33 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 cache_examples=False,
             )
 
-    # Ensure example images exist or are downloaded
-    # (Keep the existing example download logic, maybe add teacher.jpg if used in examples)
+    # Download example images with retries
     if not os.path.exists("examples"):
         os.makedirs("examples")
-        print("Created 'examples' directory. Attempting to download sample images...")
-        example_files = {
-             "astronaut.jpg": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/d1/Astronaut_-_St._Jean_Bay.jpg/640px-Astronaut_-_St._Jean_Bay.jpg",
-             "dog_bike.jpg": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/gradio/outputs_multimodal.jpg",
-             "clip_logo.png": "https://raw.githubusercontent.com/openai/CLIP/main/CLIP.png",
-             "dogs.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image8.jpg",
-             "fruits.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image9.jpg",
-             "lion.jpg": "https://huggingface.co/spaces/gradio/image-segmentation/resolve/main/images/lion.jpg",
-             "teacher.jpg": "https://images.pexels.com/photos/848117/pexels-photo-848117.jpeg?auto=compress&cs=tinysrgb&w=600" # Example with glasses/board
-        }
-        for filename, url in example_files.items():
-             filepath = os.path.join("examples", filename)
-             if not os.path.exists(filepath):
-                 try:
-                     print(f"Downloading {filename}...")
-                     wget.download(url, filepath)
-                 except Exception as e:
-                     print(f"Could not download {filename} from {url}: {e}")
-        print("Example image download attempt finished.")
-
-
-# Launch the Gradio app
+        print("Created 'examples' directory.")
+    example_files = {
+        "astronaut.jpg": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/d1/Astronaut_-_St._Jean_Bay.jpg/640px-Astronaut_-_St._Jean_Bay.jpg",
+        "dog_bike.jpg": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/gradio/outputs_multimodal.jpg",
+        "clip_logo.png": "https://raw.githubusercontent.com/openai/CLIP/main/CLIP.png",
+        "dogs.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image8.jpg",
+        "fruits.jpg": "https://raw.githubusercontent.com/ultralytics/assets/main/im/image9.jpg",
+        "lion.jpg": "https://huggingface.co/spaces/gradio/image-segmentation/resolve/main/images/lion.jpg",
+        "teacher.jpg": "https://images.pexels.com/photos/848117/pexels-photo-848117.jpeg?auto=compress&cs=tinysrgb&w=600"
+    }
+    def download_example_file(filename, url, retries=3):
+        filepath = os.path.join("examples", filename)
+        if not os.path.exists(filepath):
+            for attempt in range(retries):
+                try:
+                    print(f"Downloading {filename} (attempt {attempt + 1}/{retries})...")
+                    wget.download(url, filepath)
+                    break
+                except Exception as e:
+                    print(f"Attempt {attempt + 1} failed: {e}")
+                    if attempt + 1 == retries:
+                        print(f"Failed to download {filename} after {retries} attempts.")
+    for filename, url in example_files.items():
+        download_example_file(filename, url)
+
 if __name__ == "__main__":
-    demo.launch(debug=True) # debug=True is helpful locally
+    demo.launch(debug=True)