Update utils/utils.py

utils/utils.py

@@ -1,111 +1,111 @@
-import numpy as np
-import torch
-import lightning.pytorch as pl
-import torchmetrics
-import segmentation_models_pytorch as smp
-import skimage
-from skimage.transform import resize
-import rasterio
-from rasterio.plot import show
-import PIL
-from PIL import Image
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-def prepare_image(path:str):
-    tif = rasterio.open(path).read()[:,:,:]
-    resized = resize(tif, (9, 224, 224), 
-                order=1,
-                preserve_range=True,
-                anti_aliasing=True)
-
-    tif = resized.astype(tif.dtype)
-    tif = np.nan_to_num(tif, nan = 0) 
-    return torch.tensor(tif).type(torch.float32).unsqueeze(0)
-
-def get_gt(path:str):
-    im = np.moveaxis(rasterio.open(path).read()[[2,1,0],:,:], (0,1,2), (2,0,1))
-    return im
-
-# def make_preds_return_mask(img, model):
-#     model.eval()
-#     with torch.inference_mode():
-#         logits = model(img)
-#     pred = torch.argmax(torch.softmax(logits, dim = 1), axis = 1).to(device).numpy()
-#     return Image.fromarray(np.squeeze(np.moveaxis(pred, (0,1,2), (2,0,1)),-1)*255.0).convert("L").resize((300,224))
-
-def make_preds_return_mask(img, model):
-    model.eval()
-    with torch.inference_mode():
-        logits = model(img)
-    pred = torch.argmax(torch.softmax(logits, dim=1), axis=1).to(device).cpu().numpy()
-    mask = np.squeeze(pred)  # shape (H, W)
-    return mask  # return numpy array
-
-
-def load_model():
-    class floodLighningModel(pl.LightningModule):
-        def __init__(self, model, lr):
-            super().__init__()
-            self.model = model
-            self.lr = lr
-            # self.loss_fn = nn.CrossEntropyLoss().to(device)
-            self.loss_fn = smp.losses.DiceLoss(from_logits = True, mode = "multiclass")
-            self.iou = torchmetrics.JaccardIndex(task = "binary")
-            self.acc_fn = torchmetrics.classification.BinaryAccuracy()
-            # self.acc_fn = BinaryAccuracy()
-            self.f1_fn = torchmetrics.classification.MulticlassF1Score(num_classes = 2).to(device)
-            # self.model.save_hyperparameter(ignore = ["model"])
-
-        def forward(self, x):
-            return self.model(x)
-
-        def training_step(self, batch, batch_idx):
-            self.model.train()
-            image, gt = batch
-            logits = self.model(image)
-            loss = self.loss_fn(logits.to(device), gt.squeeze(1).type(torch.LongTensor).to(device))
-            iou = self.iou(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
-            acc = self.acc_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
-            f1 = self.f1_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
-            # f1 = self.f1_fn(torch.round(torch.sigmoid(logits)), gt)
-            self.log("loss", loss, prog_bar = True, on_step = False, on_epoch = True)
-            self.log("iou", iou, prog_bar = True, on_step = False, on_epoch = True)
-            self.log("accuracy", acc, prog_bar = True, on_step = False, on_epoch = True)
-            self.log("f1", f1, prog_bar = True, on_step = False, on_epoch = True)
-            return {"loss": loss, "f1": f1, "iou": iou, "accuracy": acc}
-
-        def validation_step(self, batch, batch_idx):
-            self.model.eval()
-            image, gt = batch
-            logits = self.model(image)
-            val_loss = self.loss_fn(logits.to(device), gt.squeeze(1).type(torch.LongTensor).to(device))
-            val_iou = self.iou(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
-            val_acc = self.acc_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
-            val_f1 = self.f1_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
-            # val_f1 = self.f1_fn(torch.round(torch.sigmoid(logits)), gt)
-            self.log("validation loss", val_loss, prog_bar = True, on_step = False, on_epoch = True)
-            self.log("validation iou", val_iou, prog_bar = True, on_step = False, on_epoch = True)
-            self.log("validation accuracy", val_acc, prog_bar = True, on_step = False, on_epoch = True)
-            self.log("validation f1", val_f1, prog_bar = True, on_step = False, on_epoch = True)
-            return {"validation loss": val_loss, "validation f1": val_f1, "validation iou": val_iou, "validation accuracy": val_acc}
-
-        def configure_optimizers(self):
-            optimizer = torch.optim.AdamW(params = self.model.parameters(), lr = self.lr)
-            return optimizer
-        
-    pl.seed_everything(2025)
-    model_ = smp.Unet(
-        encoder_name="efficientnet-b0",
-        encoder_weights="imagenet",
-        in_channels=9,
-        classes=2,
-    )
-    # model_ = smp.DPT(
-    #     encoder_name="tu-vit_base_patch16_224.augreg_in21k",
-    #     encoder_weights="imagenet",
-    #     in_channels=9,
-    #     classes=2,
-    # )
-    lightning_model = floodLighningModel.load_from_checkpoint(model = model_, lr = 5e-6, map_location=torch.device(device=device), checkpoint_path="ckpts/epoch=39-step=2760.ckpt")
+import numpy as np
+import torch
+import lightning.pytorch as pl
+import torchmetrics
+import segmentation_models_pytorch as smp
+import skimage
+from skimage.transform import resize
+import rasterio
+from rasterio.plot import show
+import PIL
+from PIL import Image
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def prepare_image(path:str):
+    tif = rasterio.open(path).read()[:,:,:]
+    resized = resize(tif, (9, 224, 224), 
+                order=1,
+                preserve_range=True,
+                anti_aliasing=True)
+
+    tif = resized.astype(tif.dtype)
+    tif = np.nan_to_num(tif, nan = 0) 
+    return torch.tensor(tif).type(torch.float32).unsqueeze(0)
+
+def get_gt(path:str):
+    im = ((np.dstack([rasterio.open(path).read(i) for i in (3,2,1)]) / np.max(rasterio.open(path).read())).clip(0,1)*255).astype(np.uint8)
+    return im
+
+# def make_preds_return_mask(img, model):
+#     model.eval()
+#     with torch.inference_mode():
+#         logits = model(img)
+#     pred = torch.argmax(torch.softmax(logits, dim = 1), axis = 1).to(device).numpy()
+#     return Image.fromarray(np.squeeze(np.moveaxis(pred, (0,1,2), (2,0,1)),-1)*255.0).convert("L").resize((300,224))
+
+def make_preds_return_mask(img, model):
+    model.eval()
+    with torch.inference_mode():
+        logits = model(img)
+    pred = torch.argmax(torch.softmax(logits, dim=1), axis=1).to(device).cpu().numpy()
+    mask = np.squeeze(pred)  # shape (H, W)
+    return mask  # return numpy array
+
+
+def load_model():
+    class floodLighningModel(pl.LightningModule):
+        def __init__(self, model, lr):
+            super().__init__()
+            self.model = model
+            self.lr = lr
+            # self.loss_fn = nn.CrossEntropyLoss().to(device)
+            self.loss_fn = smp.losses.DiceLoss(from_logits = True, mode = "multiclass")
+            self.iou = torchmetrics.JaccardIndex(task = "binary")
+            self.acc_fn = torchmetrics.classification.BinaryAccuracy()
+            # self.acc_fn = BinaryAccuracy()
+            self.f1_fn = torchmetrics.classification.MulticlassF1Score(num_classes = 2).to(device)
+            # self.model.save_hyperparameter(ignore = ["model"])
+
+        def forward(self, x):
+            return self.model(x)
+
+        def training_step(self, batch, batch_idx):
+            self.model.train()
+            image, gt = batch
+            logits = self.model(image)
+            loss = self.loss_fn(logits.to(device), gt.squeeze(1).type(torch.LongTensor).to(device))
+            iou = self.iou(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
+            acc = self.acc_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
+            f1 = self.f1_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
+            # f1 = self.f1_fn(torch.round(torch.sigmoid(logits)), gt)
+            self.log("loss", loss, prog_bar = True, on_step = False, on_epoch = True)
+            self.log("iou", iou, prog_bar = True, on_step = False, on_epoch = True)
+            self.log("accuracy", acc, prog_bar = True, on_step = False, on_epoch = True)
+            self.log("f1", f1, prog_bar = True, on_step = False, on_epoch = True)
+            return {"loss": loss, "f1": f1, "iou": iou, "accuracy": acc}
+
+        def validation_step(self, batch, batch_idx):
+            self.model.eval()
+            image, gt = batch
+            logits = self.model(image)
+            val_loss = self.loss_fn(logits.to(device), gt.squeeze(1).type(torch.LongTensor).to(device))
+            val_iou = self.iou(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
+            val_acc = self.acc_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
+            val_f1 = self.f1_fn(torch.flatten(torch.argmax(torch.softmax(logits, dim = 1), axis = 1)).to(device), torch.flatten(gt).to(device))
+            # val_f1 = self.f1_fn(torch.round(torch.sigmoid(logits)), gt)
+            self.log("validation loss", val_loss, prog_bar = True, on_step = False, on_epoch = True)
+            self.log("validation iou", val_iou, prog_bar = True, on_step = False, on_epoch = True)
+            self.log("validation accuracy", val_acc, prog_bar = True, on_step = False, on_epoch = True)
+            self.log("validation f1", val_f1, prog_bar = True, on_step = False, on_epoch = True)
+            return {"validation loss": val_loss, "validation f1": val_f1, "validation iou": val_iou, "validation accuracy": val_acc}
+
+        def configure_optimizers(self):
+            optimizer = torch.optim.AdamW(params = self.model.parameters(), lr = self.lr)
+            return optimizer
+        
+    pl.seed_everything(2025)
+    model_ = smp.Unet(
+        encoder_name="efficientnet-b0",
+        encoder_weights="imagenet",
+        in_channels=9,
+        classes=2,
+    )
+    # model_ = smp.DPT(
+    #     encoder_name="tu-vit_base_patch16_224.augreg_in21k",
+    #     encoder_weights="imagenet",
+    #     in_channels=9,
+    #     classes=2,
+    # )
+    lightning_model = floodLighningModel.load_from_checkpoint(model = model_, lr = 5e-6, map_location=torch.device(device=device), checkpoint_path="ckpts/epoch=39-step=2760.ckpt")
     return lightning_model