Add main space file

app.py

@@ -0,0 +1,75 @@
+import numpy as np
+import cv2
+import torch
+import albumentations as albu
+from pylab import imshow
+import matplotlib.pyplot as plt
+from diffusers import StableDiffusionInpaintPipeline
+from PIL import Image
+from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad
+from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image
+from cloths_segmentation.pre_trained_models import create_model
+
+# Load Cloth Segmentation Model (Ensure this is available)
+model = create_model("Unet_2020-10-30")
+model.eval()
+
+# Load Inpainting Model (Ensure this is available)
+pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
+
+def load_and_preprocess_image(image_path):
+    image = load_rgb(image_path)
+    padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)
+    transform = albu.Compose([albu.Normalize(p=1)], p=1)
+    x = transform(image=padded_image)["image"]
+    x = torch.unsqueeze(tensor_from_rgb_image(x), 0)
+    return x, image, pads
+
+def segment_cloth(image_tensor, model, pads):
+    with torch.no_grad():
+        prediction = model(image_tensor)[0][0]
+    mask = (prediction > 0).cpu().numpy().astype(np.uint8)
+    mask = unpad(mask, pads)
+    return mask
+
+def perform_inpainting(image_path, mask_path, prompt):
+    image = Image.open(image_path)
+    mask_image = Image.open(mask_path).convert("L") # Convert to single-channel grayscale
+    mask_image = mask_image.resize(image.size)  # Resize mask to match image
+
+    output_image = pipe(prompt=prompt, image=image, mask_image=mask_image).images[0]
+    return output_image
+
+def resize_and_upscale(image, new_width, new_height):
+    resized_img = cv2.resize(np.array(image), (new_width, new_height), interpolation=cv2.INTER_CUBIC)
+    return Image.fromarray(resized_img)
+
+import gradio as gr
+
+def image_segmentation_and_inpainting(image, prompt="Chinese Red and Golder Armor"):
+    x, image, pads = load_and_preprocess_image(image.name)  # Gradio provides image.name for the path
+    mask = segment_cloth(x, model, pads)
+
+    # Save mask temporarily
+    mask_path = "temp_mask.jpg"
+    plt.imsave(mask_path, mask, cmap='gray')
+
+    output_image = perform_inpainting(image.name, mask_path, prompt)
+    output_image = resize_and_upscale(output_image, 1280, 720)  # Adjust dimensions as needed
+    return output_image
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Cloth Image Segmentation and Inpainting")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Upload Image")
+            prompt_input = gr.Textbox(label="Inpainting Prompt (Optional)", value="Chinese Red and Golder Armor")
+            run_button = gr.Button("Run")
+        with gr.Column():
+            image_output = gr.Image(label="Result")
+
+    run_button.click(fn=image_segmentation_and_inpainting, inputs=[image_input, prompt_input], outputs=image_output)
+
+demo.launch(share=True) 

models/Unet_2020-10-30/weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d6e949a0e98e79fa1a814213bb73945e258699936f9352aedceda44247ab1f6
+size 53237049

requirements.txt

@@ -0,0 +1,9 @@
+torch==1.13.1
+albumentations==1.2.1
+matplotlib==3.6.2
+diffusers==0.11.1
+transformers==4.26.1
+iglovikov_helper_functions  # (Replace with the correct installation method if not on PyPI)
+cloths_segmentation         # (Replace with the correct installation method if not on PyPI)
+opencv-python-headless==4.6.0.66
+gradio==3.15.0