Create vis_model_plus_save.py

visualization/code/vis_model_plus_save.py

@@ -0,0 +1,165 @@
+from glob import glob
+import argparse
+import os
+from typing import Tuple, List
+import numpy as np
+from mmeval import MeanIoU
+from PIL import Image
+from matplotlib import pyplot as plt
+from mmseg.apis import MMSegInferencer
+from vegseg.datasets import GrassDataset
+from vegseg import models
+
+
+def get_args() -> Tuple[str, str, int]:
+    """
+    get args
+    return:
+        --device: device to use.
+        --dataset_path: dataset path.
+        --output_path: output path for saving.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", type=str, default="cuda:4")
+    parser.add_argument("--dataset_path", type=str, default="data/grass")
+    args = parser.parse_args()
+    return args.device, args.dataset_path
+
+
+def give_color_to_mask(
+    mask: Image.Image | np.ndarray, palette: List[int]
+) -> Image.Image:
+    """
+    Args:
+        mask: mask to color, numpy array or PIL Image.
+        palette: palette of dataset.
+    return:
+        mask: mask with color.
+    """
+    if isinstance(mask, np.ndarray):
+        mask = Image.fromarray(mask)
+    mask = mask.convert("P")
+    mask.putpalette(palette)
+    return mask
+
+
+def get_image_and_mask_paths(
+    dataset_path: str, num: int
+) -> Tuple[List[str], List[str]]:
+    """
+    get image and mask paths from dataset path.
+    return:
+        image_paths: list of image paths.
+        mask_paths: list of mask paths.
+    """
+    image_paths = glob(os.path.join(dataset_path, "img_dir", "*", "*.tif"))
+    if num != -1:
+        image_paths = image_paths[:num]
+    mask_paths = [
+        filename.replace("tif", "png").replace("img_dir", "ann_dir")
+        for filename in image_paths
+    ]
+    return image_paths, mask_paths
+
+
+def get_palette() -> List[int]:
+    """
+    get palette of dataset.
+    return:
+        palette: list of palette.
+    """
+    palette = []
+    palette_list = GrassDataset.METAINFO["palette"]
+    for palette_item in palette_list:
+        palette.extend(palette_item)
+    return palette
+
+
+def init_all_models(models_paths: List[str], device: str):
+    """
+    init all models
+    Args:
+        models_path (str): path to all models.
+        device (str): device to use.
+    Return:
+        models (dict): dict of models.
+    """
+    models = {}
+    for model_path in models_paths:
+        print(model_path)
+        config_path = glob(os.path.join(model_path, "*.py"))[0]
+        weight_path = glob(os.path.join(model_path, "best_mIoU_iter_*.pth"))[0]
+        inference = MMSegInferencer(
+            config_path,
+            weight_path,
+            device=device,
+            classes=GrassDataset.METAINFO["classes"],
+            palette=GrassDataset.METAINFO["palette"],
+        )
+        model_name = model_path.split(os.path.sep)[-1]
+        models[model_name] = inference
+    return models
+
+
+def main():
+    device, dataset_path = get_args()
+    image_paths, mask_paths = get_image_and_mask_paths(dataset_path, -1)
+    palette = get_palette()
+    models_paths = [
+        r"work_dirs/fcn_r50",
+        r"work_dirs/pspnet_r101",
+        r"work_dirs/deeplabv3plus_r101",
+        r"work_dirs/unet-s5-d16_deeplabv3",
+        r"work_dirs/segformer_mit-b5",
+        r"work_dirs/mask2former_swin_b",
+        r"work_dirs/dinov2_upernet",
+        r"work_dirs/experiment_p",
+    ]
+    models = init_all_models(models_paths, device)
+
+    model_order = [
+        "experiment_p",
+        "fcn_r50",
+        "pspnet_r101",
+        "deeplabv3plus_r101",
+        "unet-s5-d16_deeplabv3",
+        "segformer_mit-b5",
+        "mask2former_swin_b",
+        "dinov2_upernet"
+    ]
+
+    model_mapping = {
+        "experiment_p":"ktda",
+        "fcn_r50":"fcn",
+        "pspnet_r101":"pspnet",
+        "deeplabv3plus_r101":"deeplabv3plus",
+        "unet-s5-d16_deeplabv3":"unet",
+        "segformer_mit-b5":"segformer",
+        "mask2former_swin_b":"mask2former",
+        "dinov2_upernet":"dinov2"
+    }
+
+    os.makedirs("vis_results", exist_ok=True)
+    for model_name in model_order:
+        os.makedirs(f"data/visualization/grass/{model_name}", exist_ok=True)
+    os.makedirs(f"data/visualization/grass/input", exist_ok=True)
+    os.makedirs(f"data/visualization/grass/label", exist_ok=True)
+    for image_path, mask_path in zip(image_paths, mask_paths):
+        filename = os.path.basename(image_path)
+        for model_name, inference in models.items():
+            predictions: np.ndarray = inference(image_path)["predictions"]
+            predictions = predictions.astype(np.uint8)
+
+
+            predictions = give_color_to_mask(predictions, palette=palette)
+            
+            predictions.save(f"data/visualization/grass/{model_name}/{filename}")
+            
+
+        Image.open(image_path).save(f"data/visualization/grass/input/{filename}")
+        Image.open(mask_path).save(f"data/visualization/grass/label/{filename}")
+
+
+if __name__ == "__main__":
+    # example usage: python tools/vis_model.py --models work_dirs --device cuda:0 --dataset_path data/grass
+    main()