first working prototype

.gitignore

@@ -0,0 +1,3 @@
+venv
+__pycache__
+.DS_Store

app.py

@@ -0,0 +1,52 @@
+import gradio as gr
+from random import choice
+from lib.redraw_image import redraw_image
+from lib.find_people import find_people
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+OUTFIT_SELECTION = [
+    "Summer dress",
+    "Winter coat",
+    "Fall jacket",
+    "Formal wear",
+]
+
+
+def main():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                img_input = gr.Image(label="Image of yourself")
+                drp_outfit = gr.Dropdown(
+                    label="Select a new outfit",
+                    choices=OUTFIT_SELECTION,
+                    value=choice(OUTFIT_SELECTION),
+                )
+
+            with gr.Column():
+                btn_change = gr.Button(value="Change outfit")
+                img_output = gr.Image(label="Image of you wearing a dress")
+
+        btn_change.click(
+            generate_output, inputs=[img_input, drp_outfit], outputs=[img_output]
+        )
+
+    demo.queue().launch()
+
+
+def generate_output(img_input: PILImage, drp_outfit: str) -> PILImage:
+    img_input = Image.fromarray(img_input)
+
+    people_mask = find_people(img_input)
+    img_output = redraw_image(
+        prompt=f"person wearing {drp_outfit}",
+        image=img_input,
+        mask=people_mask,
+    )
+
+    return img_output
+
+
+if __name__ == "__main__":
+    main()

lib/cloth_seg.py

@@ -0,0 +1,30 @@
+from enum import IntEnum
+
+
+class ClothSeg(IntEnum):
+    BACKGROUND = 0
+    HAT = 1
+    HAIR = 2
+    SUNGLASSES = 3
+    UPPER_CLOTHES = 4
+    SKIRT = 5
+    PANTS = 6
+    DRESS = 7
+    BELT = 8
+    LEFT_SHOE = 9
+    RIGHT_SHOE = 10
+    FACE = 11
+    LEFT_LEG = 12
+    RIGHT_LEG = 13
+    LEFT_ARM = 14
+    RIGHT_ARM = 15
+    BAG = 16
+    SCARF = 17
+
+
+def everyhing_but_background_face_and_hair() -> list[ClothSeg]:
+    return [
+        t
+        for t in ClothSeg
+        if t not in [ClothSeg.BACKGROUND, ClothSeg.HAIR, ClothSeg.FACE]
+    ]

lib/find_people.py

@@ -0,0 +1,44 @@
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
+from PIL.Image import Image as PILImage
+from PIL import Image
+from lib.cloth_seg import everyhing_but_background_face_and_hair
+from torch import zeros_like
+from torch.nn.functional import interpolate
+from functools import cache
+
+
+def find_people(image: PILImage) -> PILImage:
+    processor = get_processor()
+    model = get_model()
+
+    inputs = processor(images=image, return_tensors="pt").to("cuda")
+    logits = model(**inputs).logits.cpu()
+
+    upsampled_logits = interpolate(
+        logits,
+        size=image.size[::-1],
+        mode="bilinear",
+        align_corners=False,
+    )
+
+    predictions = upsampled_logits.argmax(dim=1)[0]
+    mask = zeros_like(predictions)
+
+    for type in everyhing_but_background_face_and_hair():
+        mask += (predictions == type.value).long()
+
+    return Image.fromarray((mask * 255).byte().numpy(), "L")
+
+
+@cache
+def get_processor():
+    return SegformerImageProcessor.from_pretrained(
+        "mattmdjaga/segformer_b2_clothes", device="cuda"
+    )
+
+
+@cache
+def get_model():
+    return AutoModelForSemanticSegmentation.from_pretrained(
+        "mattmdjaga/segformer_b2_clothes"
+    ).to("cuda")

lib/redraw_image.py

@@ -0,0 +1,22 @@
+from PIL.Image import Image as PILImage
+from functools import cache
+from diffusers import StableDiffusionInpaintPipeline
+
+
+def redraw_image(prompt: str, image: PILImage, mask: PILImage) -> PILImage:
+    inpaint_model = get_inpaint_model()
+
+    return inpaint_model(
+        prompt=prompt,
+        image=image,
+        mask_image=mask,
+        width=image.width,
+        height=image.height,
+    ).images[0]
+
+
+@cache
+def get_inpaint_model():
+    return StableDiffusionInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting"
+    ).to("cuda")