Update app.py

app.py

@@ -1,89 +1,68 @@
 import gradio as gr
-import numpy as np
 import torch
+import numpy as np
 from diffusers import StableDiffusionInpaintPipeline
-from PIL import Image
-from segment_anything import sam_model_registry, SamPredictor
-from huggingface_hub import hf_hub_download
+from PIL import Image, ImageDraw
+from transformers import DetrImageProcessor, DetrForObjectDetection
 import spaces
 
-# Device configuration
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Step 1: Download SAM model checkpoint from Hugging Face
-checkpoint_path = hf_hub_download(
-    repo_id="Munaf1987/sam",
-    filename="sam_vit_h_4b8939.pth",
-)
-
-# Step 2: Load SAM model
-model_type = "vit_h"
-sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
-sam.to(device)
-predictor = SamPredictor(sam)
-
-# Step 3: Load Stable Diffusion Inpainting Pipeline
-sta_diff_model = "stabilityai/stable-diffusion-2-inpainting"
+# Load the Stable Diffusion Inpainting model
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    sta_diff_model,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-)
-pipe = pipe.to(device)
+    "stabilityai/stable-diffusion-2-inpainting",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
 
-# Global variables to store selected pixels and mask
-selected_pixels = []
-generated_mask = None
+# Load the DETR object detection model
+processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+detector = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50").to(device)
 
 @spaces.GPU
-def generate_mask(input_image, evt: gr.SelectData):
-    global generated_mask
-    selected_pixels.append((evt.index[0], evt.index[1]))
+def detect_and_remove(input_image, prompt):
+    if input_image is None or prompt == "":
+        return None
 
-    predictor.set_image(np.array(input_image))
-    input_points = np.array(selected_pixels)
-    input_labels = np.ones(input_points.shape[0])
+    image_np = np.array(input_image)
+    inputs = processor(images=input_image, return_tensors="pt").to(device)
 
-    masks, _, _ = predictor.predict(
-        point_coords=input_points,
-        point_labels=input_labels,
-        multimask_output=False,
-    )
+    outputs = detector(**inputs)
+    target_sizes = torch.tensor([image_np.shape[:2]]).to(device)
 
-    mask = masks[0] * 255
-    mask_image = Image.fromarray(mask.astype(np.uint8)).convert("L")
-    generated_mask = mask_image
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
 
-    return mask_image
+    mask = Image.new("L", input_image.size, 0)
+    draw = ImageDraw.Draw(mask)
 
-@spaces.GPU
-def inpaint(input_image, prompt):
-    global generated_mask
-    if input_image is None or generated_mask is None or prompt == "":
-        return None
+    # Draw boxes for "person" class only
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        if detector.config.id2label[label.item()] == "person":
+            box = [int(i) for i in box.tolist()]
+            draw.rectangle(box, fill=255)
 
-    mask_image_resized = generated_mask.resize(input_image.size)
+    if np.array(mask).sum() == 0:
+        return "No human detected."
 
+    # Inpainting
     output = pipe(
         prompt=prompt,
         image=input_image,
-        mask_image=mask_image_resized
+        mask_image=mask
     ).images[0]
 
     return output
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Stable Diffusion Inpainting with SAM Mask Selection")
+    gr.Markdown("## Automatic Human Removal and Inpainting")
 
     with gr.Row():
-        input_image = gr.Image(type="pil", label="Input Image", interactive=True)
-        mask_display = gr.Image(type="pil", label="Generated Mask")
+        input_image = gr.Image(type="pil", label="Input Image")
         output_image = gr.Image(type="pil", label="Output Image")
 
-    prompt_text = gr.Textbox(label="Prompt", placeholder="Enter a prompt for inpainting")
+    prompt_text = gr.Textbox(label="Prompt", placeholder="Example: Replace humans with cartoon background")
     submit = gr.Button("Submit")
 
-    input_image.select(generate_mask, inputs=input_image, outputs=mask_display)
-    submit.click(inpaint, inputs=[input_image, prompt_text], outputs=output_image)
+    submit.click(detect_and_remove, inputs=[input_image, prompt_text], outputs=output_image)
 
 demo.launch()