Create app.py

app.py

@@ -0,0 +1,51 @@
+import gradio as gr
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+from unet_model import UNet
+from huggingface_hub import hf_hub_download
+
+# Load trained model weights from Hugging Face Hub
+weights_path = hf_hub_download(
+    repo_id="Vizuara/unet-crack-segmentation",  # ensure this matches your repo
+    filename="unet_weights_v2.pth"                # make sure this file exists in repo
+)
+
+# Initialize and load model
+model = UNet()
+model.load_state_dict(torch.load(weights_path, map_location="cpu"))
+model.eval()
+
+# Preprocessing: same as training
+IMG_HEIGHT, IMG_WIDTH = 128, 128
+transform = transforms.Compose([
+    transforms.Resize((IMG_HEIGHT, IMG_WIDTH)),
+    transforms.ToTensor()
+])
+
+def predict(image):
+    orig_w, orig_h = image.size   # original size of uploaded image
+    img = transform(image).unsqueeze(0)  # (1,3,128,128)
+    with torch.no_grad():
+        pred = model(img)
+
+    mask = pred.squeeze(0).squeeze(0).cpu().numpy()
+    mask = (mask * 255).astype(np.uint8)  # grayscale mask
+
+    # Resize back to original size
+    mask_img = Image.fromarray(mask).resize((orig_w, orig_h), Image.NEAREST)
+    return mask_img
+
+
+# Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="UNet Crack Segmentation",
+    description="Upload a concrete surface image to get predicted crack mask"
+)
+
+if __name__ == "__main__":
+    demo.launch()