gradio update

app.py

@@ -0,0 +1,9 @@
+import gradio as gr
+from inference import get_output
+from PIL import Image
+
+def predict(image, mask):
+    return get_output(image, mask)
+
+iface = gr.Interface(fn=predict, inputs=["image", "image"], outputs="image")
+iface.launch()

handler.py

@@ -0,0 +1,38 @@
+from model import UNetInpaint
+from PIL import Image
+import torch
+import numpy as np
+import io
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model = UNetInpaint()
+        self.model.load_state_dict(torch.load("model.pth", map_location="cpu"))
+        self.model.eval()
+
+    def __call__(self, data):
+        image_bytes = data.get("image")
+        mask_bytes = data.get("mask")
+
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        mask = Image.open(io.BytesIO(mask_bytes)).convert("L")
+
+        image_np = np.array(image).astype(np.float32) / 255.0
+        mask_np = np.array(mask).astype(np.float32) / 255.0
+        mask_np = (mask_np > 0.5).astype(np.float32)
+
+        mask_np = np.expand_dims(mask_np, axis=-1)
+        image_np = np.transpose(image_np, (2, 0, 1))
+        mask_np = np.transpose(mask_np, (2, 0, 1))
+
+        image_tensor = torch.tensor(image_np) * (1 - torch.tensor(mask_np))
+        input_tensor = torch.cat([image_tensor, torch.tensor(mask_np)], dim=0).unsqueeze(0)
+
+        with torch.no_grad():
+            output = self.model(input_tensor).squeeze(0).numpy().transpose(1, 2, 0)
+            output = (np.clip(output, 0, 1) * 255).astype(np.uint8)
+            result = Image.fromarray(output)
+
+        buf = io.BytesIO()
+        result.save(buf, format="PNG")
+        return {"image": buf.getvalue()}

inference.py

@@ -0,0 +1,102 @@
+# import torch
+# import numpy as np
+# from PIL import Image
+# from model import UNetInpaint
+# import io
+
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# model = UNetInpaint(input_channels=4, output_channels=3)
+# model.load_state_dict(torch.load("inpainting_model_best.pth", map_location=device))
+# model.eval().to(device)
+
+# def preprocess(image: Image.Image, mask: Image.Image):
+#     image = image.convert("RGB").resize((256, 256))  # Resize if needed
+#     mask = mask.convert("L").resize((256, 256))
+
+#     image = np.array(image).astype(np.float32) / 255.0
+#     mask = np.array(mask).astype(np.float32) / 255.0
+#     mask = (mask > 0.5).astype(np.float32)
+#     mask = np.expand_dims(mask, axis=-1)
+
+#     image = np.transpose(image, (2, 0, 1))
+#     mask = np.transpose(mask, (2, 0, 1))
+
+#     image = torch.tensor(image, dtype=torch.float32)
+#     mask = torch.tensor(mask, dtype=torch.float32)
+
+#     image = image * (1.0 - mask)
+#     input_tensor = torch.cat([image, mask], dim=0).unsqueeze(0).to(device)
+
+#     return input_tensor
+
+# def predict(image: Image.Image, mask: Image.Image) -> Image.Image:
+#     input_tensor = preprocess(image, mask)
+
+#     with torch.no_grad():
+#         output = model(input_tensor).squeeze(0).cpu().numpy().transpose(1, 2, 0)
+#         output = np.clip(output, 0, 1)
+#         out_img = Image.fromarray((output * 255).astype(np.uint8), mode="RGB")
+#         return out_img
+
+
+# from model import UNetInpaint
+# from PIL import Image
+# import torch
+# import numpy as np
+# import io
+
+# model = UNetInpaint()
+# model.load_state_dict(torch.load("model.pth", map_location="cpu"))
+# model.eval()
+
+# def predict(image: bytes, mask: bytes) -> bytes:
+#     image = Image.open(io.BytesIO(image)).convert("RGB")
+#     mask = Image.open(io.BytesIO(mask)).convert("L")
+
+#     # preprocess
+#     image_np = np.array(image).astype(np.float32) / 255.0
+#     mask_np = np.array(mask).astype(np.float32) / 255.0
+#     mask_np = (mask_np > 0.5).astype(np.float32)
+
+#     mask_np = np.expand_dims(mask_np, axis=-1)
+#     image_np = np.transpose(image_np, (2, 0, 1))
+#     mask_np = np.transpose(mask_np, (2, 0, 1))
+
+#     image_tensor = torch.tensor(image_np) * (1 - torch.tensor(mask_np))
+#     input_tensor = torch.cat([image_tensor, torch.tensor(mask_np)], dim=0).unsqueeze(0)
+
+#     with torch.no_grad():
+#         output = model(input_tensor).squeeze(0).numpy().transpose(1, 2, 0)
+#         output = (np.clip(output, 0, 1) * 255).astype(np.uint8)
+#         result = Image.fromarray(output)
+
+#     buffer = io.BytesIO()
+#     result.save(buffer, format="PNG")
+#     return buffer.getvalue()
+
+from model import UNetInpaint
+import torch
+import numpy as np
+from PIL import Image
+
+model = UNetInpaint()
+model.load_state_dict(torch.load("model.pth", map_location="cpu"))
+model.eval()
+
+def get_output(image_pil, mask_pil):
+    image = np.array(image_pil).astype(np.float32) / 255.0
+    mask = np.array(mask_pil.convert("L")).astype(np.float32) / 255.0
+    mask = (mask > 0.5).astype(np.float32)
+
+    mask = np.expand_dims(mask, axis=-1)
+    image = np.transpose(image, (2, 0, 1))
+    mask = np.transpose(mask, (2, 0, 1))
+
+    image = torch.tensor(image) * (1 - torch.tensor(mask))
+    input_tensor = torch.cat([image, torch.tensor(mask)], dim=0).unsqueeze(0)
+
+    with torch.no_grad():
+        output = model(input_tensor).squeeze(0).numpy().transpose(1, 2, 0)
+        output = (np.clip(output, 0, 1) * 255).astype(np.uint8)
+
+    return Image.fromarray(output)

inpainting_model_best.pth

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

model.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+
+class UNetInpaint(nn.Module):
+    def __init__(self, input_channels=4, output_channels=3):
+        super().__init__()
+        self.enc1 = self.conv_block(input_channels, 64)
+        self.enc2 = self.conv_block(64, 128)
+        self.enc3 = self.conv_block(128, 256)
+        self.enc4 = self.conv_block(256, 512)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.bottleneck = self.conv_block(512, 1024)
+        self.upconv4 = self.up_conv_block(1024, 512)
+        self.dec4 = self.conv_block(1024, 512)
+        self.upconv3 = self.up_conv_block(512, 256)
+        self.dec3 = self.conv_block(512, 256)
+        self.upconv2 = self.up_conv_block(256, 128)
+        self.dec2 = self.conv_block(256, 128)
+        self.upconv1 = self.up_conv_block(128, 64)
+        self.dec1 = self.conv_block(128, 64)
+        self.out_conv = nn.Conv2d(64, output_channels, 1)
+        self.final_activation = nn.Sigmoid()
+
+    def conv_block(self, in_channels, out_channels):
+        return nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def up_conv_block(self, in_channels, out_channels):
+        return nn.Sequential(
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+        b = self.bottleneck(self.pool(e4))
+        d4 = self.upconv4(b)
+        d4 = self.dec4(torch.cat([d4, e4], dim=1))
+        d3 = self.upconv3(d4)
+        d3 = self.dec3(torch.cat([d3, e3], dim=1))
+        d2 = self.upconv2(d3)
+        d2 = self.dec2(torch.cat([d2, e2], dim=1))
+        d1 = self.upconv1(d2)
+        d1 = self.dec1(torch.cat([d1, e1], dim=1))
+        out = self.out_conv(d1)
+        return self.final_activation(out)

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+gradio
+numpy
+pillow