initial - v0

.gitignore

@@ -0,0 +1,3 @@
+.env
+__pycache__/
+flagged/

app.py

@@ -0,0 +1,29 @@
+import gradio as gr
+from PIL import Image
+from models.segmentation import SamSegmentationModel
+from models.inpainting import KandingskyInpaintingModel
+from models.product import ProductBackgroundModifier
+import torch
+
+def generate(image: Image.Image, prompt: str):
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    model = ProductBackgroundModifier(
+        segmentation_model=SamSegmentationModel(
+            model_type="vit_h", 
+            checkpoint_path="model_checkpoints/sam_vit.pth", 
+            device=device,
+        ),
+        inpainting_model=KandingskyInpaintingModel(),
+        device=device
+    )
+    generated = model.generate(image=image, prompt=prompt)
+    return generated
+
+gr.Interface(
+    fn=generate, 
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Text()
+    ], 
+    outputs=gr.Image(type="pil"),
+).launch()

model_checkpoints/sam_vit.pth

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

models/inpainting.py

@@ -0,0 +1,26 @@
+import torch
+from diffusers import AutoPipelineForInpainting
+from torchvision.transforms.functional import to_pil_image
+from PIL import Image
+
+class InpaintingModel:
+    def __init__(self) -> None:
+        pass
+    def generate(self, image: torch.Tensor, mask_image: torch.Tensor, prompt: str) -> Image.Image:
+        pass
+
+class KandingskyInpaintingModel(InpaintingModel):
+    def __init__(
+        self,
+        device = torch.device("cpu"),
+    ) -> None:
+        super().__init__()
+        self.device = device
+        self.model = AutoPipelineForInpainting.from_pretrained("kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16)
+        self.model.enable_model_cpu_offload()
+        self.negative_prompt = "deformed, ugly, disfigured"
+
+    def generate(self, image: Image.Image, mask_image: Image.Image, prompt: str) -> Image.Image:
+        output = self.model(prompt=prompt, negative_prompt=self.negative_prompt, image=image, mask_image=mask_image)
+        return output.images[0]
+

models/product.py

@@ -0,0 +1,30 @@
+import torch
+from torchvision import transforms
+from torchvision.transforms.functional import to_pil_image
+from models import segmentation, inpainting
+from PIL import Image
+
+class ProductBackgroundModifier:
+    def __init__(
+        self,
+        segmentation_model: segmentation.SegmentationModel,
+        inpainting_model: inpainting.InpaintingModel,
+        device = torch.device("cpu"),
+    ) -> None:
+        self.segmentation_model = segmentation_model
+        self.inpainting_model = inpainting_model
+        self.device = device
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Resize(1024),
+            transforms.CenterCrop((1024, 1024))
+        ])
+
+    def generate(self, image: Image.Image, prompt: str) -> Image.Image:
+        image_tensor = self.transform(image).to(self.device)
+        mask_image = self.segmentation_model.generate(image_tensor)
+        mask_image.show()
+        generated_image = self.inpainting_model.generate(image=image, mask_image=mask_image, prompt=prompt)
+        return generated_image
+
+    

models/segmentation.py

@@ -0,0 +1,43 @@
+import torch
+from torchvision.transforms.functional import to_pil_image
+from segment_anything import SamPredictor, sam_model_registry
+from PIL import Image
+
+class SegmentationModel:
+    def __init__(self) -> None:
+        pass
+    def generate(self, image: torch.Tensor) -> Image.Image:
+        pass
+
+class SamSegmentationModel(SegmentationModel):
+    def __init__(
+            self, 
+            model_type: str, 
+            checkpoint_path: str, 
+            device = torch.device("cpu"),
+        ) -> None:
+        super().__init__()
+        sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
+        sam.to(device)
+        self.device = device
+        self.model = SamPredictor(sam)
+
+    def generate(self, image: torch.Tensor) -> Image.Image:
+        _, H, W = image.size()
+        image = image.unsqueeze(0)
+        self.model.set_torch_image(image, original_image_size=(H, W))
+        center_point = [H / 2, W / 2]
+        input_point = torch.tensor([[center_point]]).to(self.device)
+        input_label = torch.tensor([[1]]).to(self.device)
+        masks, scores, logits = self.model.predict_torch(
+            point_coords=input_point,
+            point_labels=input_label,
+            boxes=None,
+            multimask_output=True
+        )
+        masks = masks.squeeze(0)
+        scores = scores.squeeze(0)
+        bmask = masks[torch.argmax(scores).item()]
+        mask_float = 1.0 - bmask.float()
+        final = torch.stack([mask_float, mask_float, mask_float])
+        return to_pil_image(final)