Create app.py

app.py

@@ -0,0 +1,93 @@
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from transformers import SamModel, SamProcessor
+
+# 1. Load the Model and Processor
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
+processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+
+def segment_object(image_data):
+    # image_data is a dictionary from the Gradio ImageEditor
+    if image_data is None or "composite" not in image_data:
+        return None
+
+    raw_image = image_data["background"].convert("RGB")
+    
+    # Get the bounding box from the editor's layers
+    # Gradio returns shapes in a list. We look for the rectangle.
+    layers = image_data.get("layers", [])
+    if not layers:
+        return raw_image # Return original if no box drawn
+
+    # For simplicity, we take the first box found
+    # In a real app, you'd iterate to find the 'crop' or 'rect' layer
+    # Here we use the composite mask logic for a beginner-friendly approach
+    
+    # Convert image for model
+    inputs = processor(raw_image, return_tensors="pt").to(device)
+    image_embeddings = model.get_image_embeddings(inputs["pixel_values"])
+
+    # In this simple version, we'll use the 'mask' drawn by the user
+    # to find the object. If you use the 'brush' or 'rect' tool:
+    mask = image_data["layers"][0].split()[-1] # Alpha channel of the drawing layer
+    mask = np.array(mask)
+    
+    # Find the coordinates of the drawn rectangle
+    coords = np.argwhere(mask > 0)
+    if coords.size == 0:
+        return raw_image
+    
+    y0, x0 = coords.min(axis=0)
+    y1, x1 = coords.max(axis=0)
+    input_boxes = [[[x0, y0, x1, y1]]]
+
+    # 2. Predict the mask
+    inputs = processor(raw_image, input_boxes=[input_boxes], return_tensors="pt").to(device)
+    inputs.pop("pixel_values", None)
+    inputs["image_embeddings"] = image_embeddings
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # 3. Process the results
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(), 
+        inputs.original_sizes.cpu(), 
+        inputs.reshaped_input_sizes.cpu()
+    )
+    
+    # Take the first mask (best guess)
+    best_mask = masks[0][0][0].numpy()
+
+    # 4. Create High-Quality White Background
+    raw_np = np.array(raw_image)
+    # Create an image where the background is white [255, 255, 255]
+    white_bg = np.ones_like(raw_np) * 255
+    
+    # Place object on white background
+    # We use the mask to choose between original pixels and white pixels
+    final_img = np.where(best_mask[..., None], raw_np, white_bg)
+    
+    return Image.fromarray(final_img.astype('uint8'))
+
+# 3. Create the Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🖌️ Object Extractor to White Background")
+    gr.Markdown("1. Upload an image. 2. Use the **Box** or **Brush** tool to highlight the object. 3. Click Submit.")
+    
+    with gr.Row():
+        input_img = gr.ImageEditor(
+            label="Input Image", 
+            type="pil", 
+            layers=True, 
+            canvas_size=(512, 512)
+        )
+        output_img = gr.Image(label="Extracted Object", type="pil")
+    
+    submit_btn = gr.Button("Extract Object")
+    submit_btn.click(segment_object, inputs=[input_img], outputs=[output_img])
+
+demo.launch()