Upload 4 files

app.py

@@ -0,0 +1,98 @@
+import os
+import torch
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+import numpy as np
+from PIL import Image
+from flask import Flask, render_template, request
+
+app = Flask(__name__)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model (assuming UNet is defined in unet.py)
+def load_model():
+    from unet import UNet
+    model = UNet().to(device)
+    model.load_state_dict(torch.load("unet_car_final.pth", map_location=device))
+    model.eval()
+    return model
+
+model = load_model()
+
+# Image transforms
+img_transform = T.Compose([
+    T.Resize((256, 256)),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225])
+])
+
+UPLOAD_FOLDER = "static"
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+
+@app.route("/", methods=["GET", "POST"])
+def index():
+    orig = "static/input.jpg" if os.path.exists(os.path.join(UPLOAD_FOLDER, "input.jpg")) else None
+    mask = "static/mask.png" if os.path.exists(os.path.join(UPLOAD_FOLDER, "mask.png")) else None
+    overlay = "static/overlay.png" if os.path.exists(os.path.join(UPLOAD_FOLDER, "overlay.png")) else None
+
+    if request.method == "POST":
+        # Handle image upload
+        if "image" in request.files:
+            file = request.files["image"]
+            if file.filename == "":
+                return "No file selected", 400
+
+            # Save uploaded image
+            img_path = os.path.join(UPLOAD_FOLDER, "input.jpg")
+            file.save(img_path)
+            orig = "static/input.jpg"
+            # Clear previous results
+            mask = None
+            overlay = None
+            for path in [os.path.join(UPLOAD_FOLDER, "mask.png"), os.path.join(UPLOAD_FOLDER, "overlay.png")]:
+                if os.path.exists(path):
+                    os.remove(path)
+
+        # Handle segmentation
+        if "segment" in request.form:
+            img_path = os.path.join(UPLOAD_FOLDER, "input.jpg")
+            if not os.path.exists(img_path):
+                return "No image available for segmentation", 400
+
+            image = Image.open(img_path).convert("RGB")
+            input_tensor = img_transform(image).unsqueeze(0).to(device)
+
+            # Predict
+            with torch.no_grad():
+                output = model(input_tensor)
+                pred_mask = torch.sigmoid(output)
+                pred_mask = (pred_mask > 0.5).float()
+
+            # Resize mask back to original image size
+            mask_resized = TF.resize(
+                TF.to_pil_image(pred_mask.squeeze().cpu()),
+                size=image.size[::-1],
+                interpolation=Image.NEAREST
+            )
+
+            # Save mask
+            mask_path = os.path.join(UPLOAD_FOLDER, "mask.png")
+            mask_resized.save(mask_path)
+
+            # Create overlay
+            mask_np = np.array(mask_resized)
+            overlay = np.array(image).copy()
+            overlay[mask_np > 128] = [255, 0, 0]
+            overlay_img = Image.fromarray(overlay)
+            overlay_path = os.path.join(UPLOAD_FOLDER, "overlay.png")
+            overlay_img.save(overlay_path)
+
+            mask = "static/mask.png"
+            overlay = "static/overlay.png"
+
+    return render_template("index.html", orig=orig, mask=mask, overlay=overlay)
+
+if __name__ == "__main__":
+    app.run(debug=True)

requirements.txt

@@ -0,0 +1,6 @@
+Flask==3.0.3
+torch==2.3.1
+torchvision==0.18.1
+pillow==10.4.0
+numpy==1.26.4
+gunicorn==22.0.0

unet.py

@@ -0,0 +1,174 @@
+import torch
+from torch import nn
+
+def center_crop(encoder_feature, target_tensor):
+    _, _, H, W = target_tensor.shape
+    enc_H, enc_W = encoder_feature.shape[2], encoder_feature.shape[3]
+    crop_H = (enc_H - H) // 2
+    crop_W = (enc_W - W) // 2
+    return encoder_feature[:, :, crop_H:crop_H+H, crop_W:crop_W+W]
+
+class UNet(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+        self.max_pool_2x2 = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        # Encoder
+        self.encoder1 = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+        )
+
+        self.encoder2 = nn.Sequential(
+            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+        )
+
+        self.encoder3 = nn.Sequential(
+            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+        )
+
+        self.encoder4 = nn.Sequential(
+            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+        )
+
+        self.encoder5 = nn.Sequential(
+            nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, padding=1),
+            nn.BatchNorm2d(1024),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=1024, out_channels=1024, kernel_size=3, padding=1),
+            nn.BatchNorm2d(1024),
+            nn.ReLU(inplace=True),
+        )
+
+
+        # Decoder
+        self.upconv_1 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=2, stride=2),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder1 = nn.Sequential(
+            nn.Conv2d(in_channels=1024, out_channels=512, kernel_size=3, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+        )
+        
+
+        self.upconv_2 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=256, out_channels=256, kernel_size=2, stride=2),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder2 = nn.Sequential(
+            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+        )
+        
+
+        self.upconv_3 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=2, stride=2),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder3 = nn.Sequential(
+            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+        )
+        
+
+        self.upconv_4 = nn.Sequential(
+            nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=2, stride=2),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder4 = nn.Sequential(
+            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+        )
+
+        self.final_conv = nn.Conv2d(64, 1, kernel_size=1)
+
+
+    def forward(self, image):
+
+        # Encoder
+        e1 = self.encoder1(image)
+        e2 = self.encoder2(self.max_pool_2x2(e1))
+        e3 = self.encoder3(self.max_pool_2x2(e2))
+        e4 = self.encoder4(self.max_pool_2x2(e3))
+        e5 = self.encoder5(self.max_pool_2x2(e4))
+
+        # Decoder
+        d1 = self.upconv_1(e5)
+        c1 = center_crop(e4, d1)
+        d1 = torch.cat([d1, c1], dim=1)
+        d1 = self.decoder1(d1)
+        
+        d2 = self.upconv_2(d1)
+        c2 = center_crop(e3, d2)
+        d2 = torch.cat([d2, c2], dim=1)
+        d2 = self.decoder2(d2)
+
+        d3 = self.upconv_3(d2)
+        c3 = center_crop(e2, d3)
+        d3 = torch.cat([d3, c3], dim=1)
+        d3 = self.decoder3(d3)
+
+        d4 = self.upconv_4(d3)
+        c4 = center_crop(e1, d4)
+        d4 = torch.cat([d4, c4], dim=1)
+        d4 = self.decoder4(d4)
+
+        # Final output
+        out = self.final_conv(d4)
+
+        return out

unet_car_final.pth

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