Create app.py

app.py

@@ -0,0 +1,158 @@
+import streamlit as st
+import requests
+import zipfile
+import io
+import os
+import sys
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import cv2
+from PIL import Image
+from diffusers import StableDiffusionInpaintPipeline, EulerDiscreteScheduler
+import copy
+
+# ========== Download SAM Repo ==========
+def download_sam_repo():
+    repo_url = "https://github.com/facebookresearch/segment-anything/archive/refs/heads/main.zip"
+    repo_dir = "segment-anything-main"
+    if not os.path.exists(repo_dir):
+        st.info("🔽 Downloading Segment Anything repo...")
+        response = requests.get(repo_url)
+        if response.status_code == 200:
+            with zipfile.ZipFile(io.BytesIO(response.content)) as zip_ref:
+                zip_ref.extractall(".")
+            st.success("✅ Segment Anything repo downloaded!")
+        else:
+            st.error(f"❌ Failed to download repo: {response.status_code}")
+
+download_sam_repo()
+sys.path.append(os.path.abspath("segment-anything-main"))
+
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+
+# ========== Download SAM Model Checkpoint ==========
+def download_file(url, output_path):
+    if not os.path.exists(output_path):
+        st.info(f"🔽 Downloading {os.path.basename(output_path)}...")
+        response = requests.get(url, stream=True)
+        with open(output_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+        st.success(f"✅ Downloaded {os.path.basename(output_path)}")
+
+sam_url = "https://huggingface.co/camenduru/segment_anything/resolve/main/sam_vit_h_4b8939.pth"
+download_file(sam_url, "sam_vit_h_4b8939.pth")
+
+# ========== Load SAM Model ==========
+@st.cache_resource
+def load_sam_model():
+    sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
+    sam.to("cuda")
+    mask_generator = SamAutomaticMaskGenerator(
+        model=sam,
+        points_per_side=32,
+        pred_iou_thresh=0.99,
+        stability_score_thresh=0.92,
+        crop_n_layers=1,
+        crop_n_points_downscale_factor=2,
+        min_mask_region_area=100,
+    )
+    return mask_generator
+
+# ========== Load SD Model ==========
+@st.cache_resource
+def load_sd_pipeline():
+    model_dir = 'stabilityai/stable-diffusion-2-inpainting'
+    scheduler = EulerDiscreteScheduler.from_pretrained(model_dir, subfolder='scheduler')
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        model_dir,
+        scheduler=scheduler,
+        torch_dtype=torch.float16,
+        revision="fp16"
+    ).to("cuda")
+    pipe.enable_xformers_memory_efficient_attention()
+    return pipe
+
+# ========== Display masks ==========
+def show_masks(image, masks):
+    fig, ax = plt.subplots(figsize=(10, 10))
+    ax.imshow(image)
+    for i, mask in enumerate(masks):
+        m = mask['segmentation']
+        color = np.random.random((1, 3)).tolist()[0]
+        img = np.ones((m.shape[0], m.shape[1], 3))
+        for j in range(3):
+            img[:, :, j] = color[j]
+        ax.imshow(np.dstack((img, m * 0.35)))
+        contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if contours:
+            cnt = contours[0]
+            M = cv2.moments(cnt)
+            if M["m00"] != 0:
+                cx = int(M["m10"] / M["m00"])
+                cy = int(M["m01"] / M["m00"])
+                ax.text(cx, cy, str(i), color='white', fontsize=16, ha='center', va='center', fontweight='bold')
+    ax.axis('off')
+    st.pyplot(fig)
+
+# ========== Image Grid ==========
+def create_image_grid(original_image, images, names, rows, columns):
+    images = copy.copy(images)
+    names = copy.copy(names)
+
+    images.insert(0, original_image)
+    names.insert(0, "Original")
+
+    fig, axes = plt.subplots(rows, columns, figsize=(15, 15))
+    for idx, (img, name) in enumerate(zip(images, names)):
+        row, col = divmod(idx, columns)
+        axes[row, col].imshow(img)
+        axes[row, col].set_title(name)
+        axes[row, col].axis('off')
+    for idx in range(len(images), rows * columns):
+        row, col = divmod(idx, columns)
+        axes[row, col].axis('off')
+    plt.tight_layout()
+    st.pyplot(fig)
+
+# ========== Streamlit UI ==========
+st.title("🎨 Segment & Inpaint with Streamlit")
+
+uploaded_file = st.file_uploader("Upload an Image", type=['png', 'jpg', 'jpeg'])
+
+if uploaded_file:
+    source_image = Image.open(uploaded_file).convert("RGB").resize((512, 512))
+    st.image(source_image, caption="Uploaded Image", use_column_width=True)
+
+    mask_gen = load_sam_model()
+    masks = mask_gen.generate(np.asarray(source_image))
+    st.write(f"Number of Segments Found: {len(masks)}")
+
+    show_masks(source_image, masks)
+
+    mask_idx = st.number_input(f"Select Mask Index (0 to {len(masks)-1})", min_value=0, max_value=len(masks)-1, value=0)
+    prompt = st.text_input("Enter Inpainting Prompt", "a skirt full of text")
+    generate = st.button("Generate Inpainting")
+
+    if generate:
+        segmentation_mask = masks[mask_idx]['segmentation']
+        stable_mask = Image.fromarray(segmentation_mask * 255).convert("RGB")
+
+        pipe = load_sd_pipeline()
+        generator = torch.Generator(device="cuda").manual_seed(77)
+
+        images = []
+        for i in range(4):
+            result = pipe(
+                prompt=prompt,
+                guidance_scale=7.5,
+                num_inference_steps=50,
+                generator=generator,
+                image=source_image,
+                mask_image=stable_mask
+            ).images[0]
+            images.append(result)
+
+        create_image_grid(source_image, images, [prompt]*4, 2, 3)