Clean code

app.py

@@ -4,6 +4,7 @@ import torch.nn.functional as F
 from torchvision import transforms
 import segmentation_models_pytorch as smp
 import gradio as gr
+import numpy as np
 import os
 
 config = {
@@ -31,7 +32,9 @@ class_rgb_colors = [
 ]
 
 
-def label_to_onehot(mask, num_classes):
+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),
@@ -50,43 +53,36 @@ model.to(device)
 model.eval()
 
 
-# mean = [0.4085, 0.3798, 0.2822]
-# std = [0.1410, 0.1051, 0.0927]
 # transforms
 downsize_t = transforms.Resize(config["downsize_res"], antialias=True)
 transform = transforms.Compose(
     [
         transforms.ToTensor(),
-        # transforms.Normalize(mean, std),
     ]
 )
 
 
-def get_overlay(sat_img, preds, alpha):
-    class_rgb_colors = [
-        (0, 255, 255),
-        (255, 255, 0),
-        (255, 0, 255),
-        (0, 255, 0),
-        (0, 0, 255),
-        (255, 255, 255),
-        (0, 0, 0),
-    ]
-    masks = preds.squeeze()
-    masks = label_to_onehot(masks, 7)
+def get_overlay(image: torch.Tensor, preds: torch.Tensor, alpha: float) -> torch.Tensor:
+    """Generates the segmentation ovelay for an satellite image"""
+
+    masks = preds.squeeze().label_to_onehot(masks, 7)
     overlay = draw_segmentation_masks(
-        sat_img, masks=masks, alpha=alpha, colors=class_rgb_colors
+        image, masks=masks, alpha=alpha, colors=class_rgb_colors
     )
     return overlay
 
 
-def segment(numpy_arr):
-    sat_img_arr = torch.from_numpy(numpy_arr)
-    sat_img_arr = torch.permute(sat_img_arr, (2, 0, 1))
-    sat_img_pil = transforms.functional.to_pil_image(sat_img_arr)
+def hwc_to_chw(image_tensor: torch.Tensor) -> torch.Tensor:
+    return torch.permute(image_tensor, (2, 0, 1))
+
+
+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(sat_img_pil).unsqueeze(0)
+    X = transform(pil_image).unsqueeze(0)
     X = X.to(device)
     X_down = downsize_t(X)
     # forward pass
@@ -94,26 +90,23 @@ def segment(numpy_arr):
     preds = torch.argmax(logits, 1).detach()
     # 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(
+        get_overlay(torch.zeros_like(image_tensor), preds, 1)
+    ).numpy()
 
-    # ger rbg formatted images
-    overlay = get_overlay(sat_img_arr, preds, 0.2)
-    raw_masks = get_overlay(torch.zeros_like(sat_img_arr), preds, 1)
-    raw_masks = torch.permute(raw_masks, (1, 2, 0))
-    overlay = torch.permute(overlay, (1, 2, 0))
-
-    return raw_masks.numpy(), overlay.numpy()
+    return raw_segmentation, segmentation_overlay
 
 
 i = gr.inputs.Image()
 o = [gr.Image(), gr.Image()]
 images_dir = "sample_sat_images/"
-
-image_ids = os.listdir(images_dir)
-examples = [f"{images_dir}/{image_id}" for image_id in image_ids]
+examples = [f"{images_dir}/{image_id}" for image_id in os.listdir(images_dir)]
 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
+    segment, i, o, examples=examples, title=title, description=description, cache_examples=True
 )
-iface.launch(debug=True)
+iface.launch()