Create app.py

app.py

@@ -0,0 +1,106 @@
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import StableDiffusionInpaintPipeline
+from segment_anything import sam_model_registry, SamPredictor
+from groundingdino.util.inference import load_model, load_image, predict, annotate
+
+# Device configuration
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load Grounding DINO (human detection)
+grounding_model = load_model("ShilongLiu/GroundingDINO-SwinB")  # Public Hugging Face model
+
+# Load SAM model
+sam_checkpoint = "facebook/sam-vit-huge"
+sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
+sam.to(device)
+predictor = SamPredictor(sam)
+
+# Load Stable Diffusion Inpainting Pipeline
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-inpainting",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
+pipe = pipe.to(device)
+
+def detect_and_segment(input_image, prompt):
+    # Convert image to numpy
+    image_np = np.array(input_image)
+    predictor.set_image(image_np)
+
+    # Grounding DINO detection
+    boxes, logits, phrases = predict(
+        model=grounding_model,
+        image=input_image,
+        caption=prompt,
+        box_threshold=0.35,
+        text_threshold=0.25
+    )
+
+    if len(boxes) == 0:
+        return None, None, "No objects detected."
+
+    # Prepare mask
+    transformed_boxes = boxes * torch.tensor([input_image.width, input_image.height, input_image.width, input_image.height])
+    transformed_boxes = transformed_boxes.cpu().numpy()
+
+    input_points = []
+    input_labels = []
+    for box in transformed_boxes:
+        x_center = int((box[0] + box[2]) / 2)
+        y_center = int((box[1] + box[3]) / 2)
+        input_points.append([x_center, y_center])
+        input_labels.append(1)
+
+    masks, _, _ = predictor.predict(
+        point_coords=np.array(input_points),
+        point_labels=np.array(input_labels),
+        multimask_output=False,
+    )
+
+    final_mask = np.zeros_like(masks[0])
+    for mask in masks:
+        final_mask = np.logical_or(final_mask, mask)
+
+    final_mask = (final_mask * 255).astype(np.uint8)
+    mask_image = Image.fromarray(final_mask).convert("L")
+    return input_image, mask_image, "Mask generated successfully."
+
+def inpaint(input_image, mask_image, inpaint_prompt):
+    if input_image is None or mask_image is None or inpaint_prompt == "":
+        return None
+
+    image_resized = input_image.resize((512, 512))
+    mask_resized = mask_image.resize((512, 512))
+
+    output = pipe(
+        prompt=inpaint_prompt,
+        image=image_resized,
+        mask_image=mask_resized
+    ).images[0]
+
+    # Resize back to original
+    output = output.resize(input_image.size)
+    return output
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## Remove Humans and Replace with Cartoon / Imaginary Characters")
+
+    with gr.Row():
+        input_image = gr.Image(type="pil", label="Upload Image")
+        mask_display = gr.Image(type="pil", label="Generated Mask")
+        output_image = gr.Image(type="pil", label="Final Output")
+
+    detect_prompt = gr.Textbox(label="Detection Prompt", value="human", placeholder="What objects to detect? (e.g., human)")
+    inpaint_prompt = gr.Textbox(label="Inpainting Prompt", placeholder="What to replace with? (e.g., cartoon human, anime boy)")
+
+    detect_button = gr.Button("Detect and Generate Mask")
+    inpaint_button = gr.Button("Inpaint with Replacement")
+
+    detect_button.click(fn=detect_and_segment, inputs=[input_image, detect_prompt], outputs=[input_image, mask_display, gr.Textbox(label="Status")])
+    inpaint_button.click(fn=inpaint, inputs=[input_image, mask_display, inpaint_prompt], outputs=output_image)
+
+demo.launch()