small bug fix

app.py

@@ -22,7 +22,7 @@ config = {
     },
 }
 
-class_rgb_colors = [
+colors = [
     (0, 255, 255),
     (255, 255, 0),
     (255, 0, 255),
@@ -33,16 +33,6 @@ class_rgb_colors = [
 ]
 
 
-def label_to_onehot(mask: torch.Tensor, num_classes: int) -> torch.Tensor:
-    """Transforms a tensor from label encoding to one hot encoding in boolean dtype"""
-
-    dims_p = (2, 0, 1) if mask.ndim == 2 else (0, 3, 1, 2)
-    return torch.permute(
-        F.one_hot(mask.type(torch.long), num_classes=num_classes).type(torch.bool),
-        dims_p,
-    )
-
-
 cp_path = "CP_epoch20.pth"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -63,13 +53,21 @@ transform = transforms.Compose(
 )
 
 
+def label_to_onehot(mask: torch.Tensor, num_classes: int) -> torch.Tensor:
+    """Transforms a tensor from label encoding to one hot encoding in boolean dtype"""
+
+    dims_p = (2, 0, 1) if mask.ndim == 2 else (0, 3, 1, 2)
+    return torch.permute(
+        F.one_hot(mask.type(torch.long), num_classes=num_classes).type(torch.bool),
+        dims_p,
+    )
+
+
 def get_overlay(image: torch.Tensor, preds: torch.Tensor, alpha: float) -> torch.Tensor:
     """Generates the segmentation ovelay for an satellite image"""
 
     masks = label_to_onehot(preds.squeeze(), 7)
-    overlay = draw_segmentation_masks(
-        image, masks=masks, alpha=alpha, colors=class_rgb_colors
-    )
+    overlay = draw_segmentation_masks(image, masks=masks, alpha=alpha, colors=colors)
     return overlay
 
 
@@ -77,11 +75,14 @@ def hwc_to_chw(image_tensor: torch.Tensor) -> torch.Tensor:
     return torch.permute(image_tensor, (2, 0, 1))
 
 
+def chw_to_hwc(image_tensor: torch.Tensor) -> torch.Tensor:
+    return torch.permute(image_tensor, (1, 2, 0))
+
+
 def segment(satellite_image: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     image_tensor = torch.from_numpy(satellite_image)
     image_tensor = hwc_to_chw(image_tensor)
     pil_image = transforms.functional.to_pil_image(image_tensor)
-
     # preprocess image
     X = transform(pil_image).unsqueeze(0)
     X = X.to(device)
@@ -92,8 +93,8 @@ def segment(satellite_image: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     # resize to evaluate with the original image
     preds = transforms.functional.resize(preds, X.shape[-2:], antialias=True)
     # get rbg formatted images
-    segmentation_overlay = hwc_to_chw(get_overlay(image_tensor, preds, 0.2)).numpy()
-    raw_segmentation = hwc_to_chw(
+    segmentation_overlay = chw_to_hwc(get_overlay(image_tensor, preds, 0.2)).numpy()
+    raw_segmentation = chw_to_hwc(
         get_overlay(torch.zeros_like(image_tensor), preds, 1)
     ).numpy()
 
@@ -108,6 +109,12 @@ title = "Satellite Images Landcover Segmentation"
 description = "Upload an image or select from examples to segment"
 
 iface = gr.Interface(
-    segment, i, o, examples=examples, title=title, description=description, cache_examples=True
+    segment,
+    i,
+    o,
+    examples=examples,
+    title=title,
+    description=description,
+    cache_examples=True,
 )
 iface.launch()