Merge pull request #184 from ENSTA-U2IS-AI/dev

🔨 Rework FrostImages dataset & Rename a plotting function

Author: o-laurent

auto_tutorial_source/Bayesian_Methods/tutorial_bayesian.py

@@ -119,7 +119,7 @@
 #
 # Now that the model is trained, let's test it on MNIST.
 # Please note that we apply a reshape to the logits to determine the dimension corresponding to the ensemble
-# and to the batch. As for TorchUncertainty 0.5.1, the ensemble dimension is merged with the batch dimension
+# and to the batch. As for TorchUncertainty 0.5.2, the ensemble dimension is merged with the batch dimension
 # in this order (num_estimator x batch, classes).
 
 import matplotlib.pyplot as plt

docs/source/conf.py

@@ -20,7 +20,7 @@
     f"{datetime.now().year!s}, Adrien Lafage and Olivier Laurent"
 )
 author = "Adrien Lafage and Olivier Laurent"
-release = "0.5.1"
+release = "0.5.2"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "torch_uncertainty"
-version = "0.5.1"
+version = "0.5.2"
 authors = [
     { name = "ENSTA U2IS", email = "olivier.laurent@ensta-paris.fr" },
     { name = "Adrien Lafage", email = "adrienlafage@outlook.com" },
@@ -41,7 +41,7 @@ dependencies = [
 
 [project.optional-dependencies]
 experiments = ["tensorboard", "huggingface-hub>=0.31", "safetensors"]
-image = ["kornia", "h5py", "opencv-python"]
+image = ["kornia", "h5py", "opencv-python", "torch-uncertainty-assets"]
 tabular = ["pandas"]
 dev = [
     "torch_uncertainty[experiments,image]",

torch_uncertainty/datasets/frost.py

@@ -1,78 +1,43 @@
-import logging
 from collections.abc import Callable
+from importlib import util
+from importlib.abc import Traversable
+from importlib.resources import files
 from pathlib import Path
 from typing import Any
 
 from PIL import Image
 from torchvision.datasets import VisionDataset
-from torchvision.datasets.utils import (
-    check_integrity,
-    download_and_extract_archive,
-)
 
+FROST_ASSETS_MOD = "torch_uncertainty_assets.frost"
+tu_assets_installed = util.find_spec("torch_uncertainty_assets")
 
-def pil_loader(path: Path) -> Image.Image:
-    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+
+def pil_loader(path: Path | Traversable) -> Image.Image:
     with path.open("rb") as f:
         img = Image.open(f)
         return img.convert("RGB")
 
 
-class FrostImages(VisionDataset):  # TODO: Use ImageFolder
-    url = "https://zenodo.org/records/10438904/files/frost.zip"
-    zip_md5 = "d82f29f620d43a68e71e34b28f7c35cb"
-    filename = "frost.zip"
-    samples = [
-        "frost1.png",
-        "frost2.png",
-        "frost3.jpg",
-        "frost4.jpg",
-        "frost5.jpg",
-    ]
-
+class FrostImages(VisionDataset):
     def __init__(
         self,
-        root: str | Path,
-        transform: Callable[..., Any] | None,
+        transform: Callable[..., Any] | None = None,
         target_transform: Callable[..., Any] | None = None,
-        download: bool = False,
     ) -> None:
-        self.root = Path(root)
-
-        if download:
-            self.download()
-
-        if not self._check_integrity():
-            raise RuntimeError(
-                "Dataset not found or corrupted. You can use download=True to download it."
+        if not tu_assets_installed:  # coverage: ignore
+            raise ImportError(
+                "The torch-uncertainty-assets library is not installed. Please install"
+                "torch_uncertainty with the image option:"
+                """pip install -U "torch_uncertainty[image]"."""
             )
-
         super().__init__(
-            self.root / "frost",
+            FROST_ASSETS_MOD,
             transform=transform,
             target_transform=target_transform,
         )
         self.loader = pil_loader
-
-    def _check_integrity(self) -> bool:
-        fpath = self.root / self.filename
-        return check_integrity(
-            fpath,
-            self.zip_md5,
-        )
-
-    def download(self) -> None:
-        if self._check_integrity():
-            logging.info("Files already downloaded and verified")
-            return
-
-        download_and_extract_archive(
-            self.url,
-            download_root=self.root,
-            filename=self.filename,
-            md5=self.zip_md5,
-        )
-        logging.info("Downloaded %s to %s.", self.filename, self.root)
+        sample_path = files(FROST_ASSETS_MOD)
+        self.samples = [sample_path.joinpath(f"frost{i}.jpg") for i in range(1, 6)]
 
     def __getitem__(self, index: int) -> Any:
         """Get the samples of the dataset.
@@ -83,8 +48,7 @@ def __getitem__(self, index: int) -> Any:
         Returns:
             tuple: (sample, target) where target is class_index of the target class.
         """
-        path = self.root / self.samples[index]
-        sample = self.loader(path)
+        sample = self.loader(self.samples[index])
         if self.transform is not None:
             sample = self.transform(sample)
         return sample

torch_uncertainty/routines/segmentation.py

@@ -37,7 +37,7 @@
 )
 from torch_uncertainty.post_processing import PostProcessing
 from torch_uncertainty.utils import csv_writer
-from torch_uncertainty.utils.plotting import show
+from torch_uncertainty.utils.plotting import show_segmentation_predictions
 
 
 class SegmentationRoutine(LightningModule):
@@ -421,7 +421,7 @@ def log_segmentation_plots(self) -> None:
 
             self.logger.experiment.add_figure(
                 f"Segmentation results/{i}",
-                show(pred_mask, gt_mask),
+                show_segmentation_predictions(pred_mask, gt_mask),
             )
 
     def subsample(self, pred: Tensor, target: Tensor) -> tuple[Tensor, Tensor]:

torch_uncertainty/transforms/corruption.py

@@ -584,7 +584,7 @@ def __init__(self, severity: int, seed: int | None = None) -> None:
         super().__init__(severity)
         self.rng = np.random.default_rng(seed)
         self.mix = [(1, 0.4), (0.8, 0.6), (0.7, 0.7), (0.65, 0.7), (0.6, 0.75)][severity - 1]
-        self.frost_ds = FrostImages("./data", download=True, transform=ToTensor())
+        self.frost_ds = FrostImages(transform=ToTensor())
 
     def forward(self, img: Tensor) -> Tensor:
         if self.severity == 0:

torch_uncertainty/utils/__init__.py

@@ -5,6 +5,6 @@
 from .evaluation_loop import TUEvaluationLoop
 from .hub import load_hf
 from .misc import csv_writer
-from .plotting import plot_hist, show
+from .plotting import plot_hist, show_segmentation_predictions
 from .trainer import TUTrainer
 from .transforms import interpolation_modes_from_str

torch_uncertainty/utils/plotting.py

@@ -7,9 +7,9 @@
 from torch import Tensor
 
 
-def show(prediction: Tensor, target: Tensor) -> Figure:
+def show_segmentation_predictions(prediction: Tensor, target: Tensor) -> Figure:
     imgs = [prediction, target]
-    fig, axs = plt.subplots(ncols=len(imgs), figsize=(12, 6))
+    fig, axs = plt.subplots(ncols=2, figsize=(12, 6), dpi=300)
     for i, img in enumerate(imgs):
         img = img.detach()
         img = F.to_pil_image(img)
@@ -18,7 +18,6 @@ def show(prediction: Tensor, target: Tensor) -> Figure:
 
     axs[0].set(title="Prediction")
     axs[1].set(title="Ground Truth")
-
     return fig