Update app.py

app.py

@@ -25,42 +25,16 @@ transform = transforms.Compose([
 ])
 
 def predict(image):
-    # Preprocess
     orig_w, orig_h = image.size
-    img = transform(image).unsqueeze(0)  # shape [1,3,H,W] (float tensor)
-
+    img = transform(image).unsqueeze(0).float()
     with torch.no_grad():
-        pred = model(img)   # shape: [1,C,H,W] where C>=1
-
-    pred = pred.cpu()  # move to cpu for numpy ops
-
-    # If single-channel output (binary segmentation)
-    if pred.shape[1] == 1:
-        # Use sigmoid -> probability map -> binary mask
-        prob = torch.sigmoid(pred)       # [1,1,H,W]
-        prob_np = prob.squeeze(0).squeeze(0).numpy()   # [H,W]
-        # optional: threshold (0.5) to binary mask
-        mask = (prob_np >= 0.5).astype(np.uint8) * 255
-
-        mask_img = Image.fromarray(mask).resize((orig_w, orig_h), Image.NEAREST)
-        return mask_img.convert("L")   # grayscale
-
-    # If multi-channel output (C-class segmentation)
-    else:
-        # Convert logits to class map with argmax
-        class_map = pred.argmax(dim=1).squeeze(0).numpy().astype(np.uint8)  # [H,W], values 0..C-1
-
-        # If you want a greyscale visualization: scale classes into 0..255
-        if pred.shape[1] <= 256:
-            # map class indices -> greyscale values (simple linear mapping)
-            vis = (class_map.astype(np.float32) / (pred.shape[1]-1) * 255.0).astype(np.uint8)
-            mask_img = Image.fromarray(vis).resize((orig_w, orig_h), Image.NEAREST)
-            return mask_img.convert("L")
-        else:
-            # If >256 classes (rare), just return raw indices scaled modulo 256
-            vis = (class_map % 256).astype(np.uint8)
-            mask_img = Image.fromarray(vis).resize((orig_w, orig_h), Image.NEAREST)
-            return mask_img.convert("L")
+        pred = model(img)          # [1,7,H,W]
+    # convert to class map
+    class_map = pred.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.uint8)  # H,W
+    # if you want a binary crack mask from classes (treat class 0 as background)
+    binary_mask = (class_map != 0).astype(np.uint8) * 255
+    # return grayscale mask (or return class_map visual)
+    return Image.fromarray(binary_mask).resize((orig_w, orig_h), Image.NEAREST).convert("L")
 
 
 # Gradio interface