endoreg-db 0.8.9.32__py3-none-any.whl

endoreg_db/utils/ai/multilabel_classification_net.py

@@ -0,0 +1,270 @@
+import logging
+from pathlib import Path
+
+import torch
+from torchvision import models
+import torch.nn as nn
+from pytorch_lightning import LightningModule
+import numpy as np
+from safetensors.torch import load_file
+from sklearn.metrics import precision_score, recall_score, f1_score
+
+try:  # Torchvision >= 0.13 exposes explicit weight enums
+    from torchvision.models import EfficientNet_B4_Weights, RegNet_X_800MF_Weights
+except ImportError:  # pragma: no cover - compatibility with older torchvision
+    EfficientNet_B4_Weights = None
+    RegNet_X_800MF_Weights = None
+
+logger = logging.getLogger(__name__)
+
+METRICS_ON_STEP = False
+
+
+def calculate_metrics(pred, target, threshold=0.5):
+    pred = np.array(pred > threshold, dtype=float)
+    return {
+        "micro/precision": precision_score(
+            y_true=target, y_pred=pred, average="micro", zero_division=0
+        ),
+        "micro/recall": recall_score(
+            y_true=target, y_pred=pred, average="micro", zero_division=0
+        ),
+        "micro/f1": f1_score(
+            y_true=target, y_pred=pred, average="micro", zero_division=0
+        ),
+        "macro/precision": precision_score(
+            y_true=target, y_pred=pred, average="macro", zero_division=0
+        ),
+        "macro/recall": recall_score(
+            y_true=target, y_pred=pred, average="macro", zero_division=0
+        ),
+        "macro/f1": f1_score(
+            y_true=target, y_pred=pred, average="macro", zero_division=0
+        ),
+        "samples/precision": precision_score(
+            y_true=target, y_pred=pred, average=None, zero_division=0
+        ),
+        "samples/recall": recall_score(
+            y_true=target, y_pred=pred, average=None, zero_division=0
+        ),
+        "samples/f1": f1_score(
+            y_true=target, y_pred=pred, average=None, zero_division=0
+        ),
+    }
+
+
+def _load_torchvision_backbone(factory, *, weights_enum=None, load_pretrained=False):
+    """Instantiate a torchvision model without triggering unwanted downloads."""
+    if weights_enum is not None:
+        try:
+            weights = weights_enum.DEFAULT if load_pretrained else None
+            return factory(weights=weights)
+        except (TypeError, AttributeError):
+            # Fall back to legacy keyword on older torchvision versions
+            pass
+
+    try:
+        return factory(pretrained=load_pretrained)
+    except TypeError:
+        # Newer torchvision versions removed the pretrained kwarg; call without hints
+        try:
+            return factory()
+        except Exception as exc:  # pragma: no cover - surfaced to caller for visibility
+            raise RuntimeError(
+                "Failed to instantiate torchvision backbone with load_pretrained="
+                f"{load_pretrained}."
+            ) from exc
+
+
+class MultiLabelClassificationNet(LightningModule):
+    def __init__(
+        self,
+        labels=None,
+        lr=6e-3,
+        weight_decay=0.001,
+        pos_weight=2,
+        model_type="EfficientNetB4",
+        load_imagenet_weights: bool = False,
+        track_hparams: bool = True,
+    ):
+        super().__init__()
+        if track_hparams:
+            self.save_hyperparameters()
+        if labels is None:
+            raise ValueError(
+                "labels must be provided to initialize MultiLabelClassificationNet"
+            )
+
+        self.model_type = model_type
+        self.labels = list(labels)
+        self.n_classes = len(self.labels)
+        self.val_preds: list[np.ndarray] = []
+        self.val_targets: list[np.ndarray] = []
+        self.pos_weight = pos_weight
+        self.weight_decay = weight_decay
+        self.lr = lr
+        self.sigm = nn.Sigmoid()
+
+        if model_type == "EfficientNetB4":
+            self.model = _load_torchvision_backbone(
+                models.efficientnet_b4,
+                weights_enum=EfficientNet_B4_Weights,
+                load_pretrained=load_imagenet_weights,
+            )
+            num_ftrs = self.model.classifier[1].in_features
+            self.model.classifier[1] = nn.Linear(num_ftrs, len(labels))
+
+        elif model_type == "RegNetX800MF":
+            self.model = _load_torchvision_backbone(
+                models.regnet_x_800mf,
+                weights_enum=RegNet_X_800MF_Weights,
+                load_pretrained=load_imagenet_weights,
+            )
+            num_ftrs = self.model.fc.in_features
+            self.model.fc = nn.Linear(num_ftrs, len(labels))
+
+        self.criterion = nn.BCEWithLogitsLoss(
+            pos_weight=torch.Tensor([self.pos_weight] * len(self.labels))
+        )
+
+    @classmethod
+    def load_from_checkpoint(cls, checkpoint_path, *args, **kwargs):
+        path = Path(checkpoint_path)
+        suffix = path.suffix.lower()
+
+        if suffix == ".safetensors":
+            map_location = kwargs.pop("map_location", "cpu")
+            strict = kwargs.pop("strict", True)
+            labels = kwargs.pop("labels", None)
+            if not labels:
+                raise ValueError(
+                    "labels must be provided when loading .safetensors checkpoints"
+                )
+            model_type = kwargs.pop("model_type", None) or "EfficientNetB4"
+            load_imagenet = kwargs.pop("load_imagenet_weights", False)
+
+            device = (
+                torch.device(map_location)
+                if map_location is not None
+                else torch.device("cpu")
+            )
+            if isinstance(device, torch.device):
+                device_hint = (
+                    f"{device.type}:{device.index}"
+                    if device.index is not None
+                    else device.type
+                )
+            else:
+                device_hint = device
+
+            state_dict = load_file(path, device=device_hint)
+
+            instance = cls(
+                labels=labels,
+                model_type=model_type,
+                load_imagenet_weights=load_imagenet,
+                track_hparams=False,
+                *args,
+                **kwargs,
+            )
+            missing, unexpected = instance.load_state_dict(state_dict, strict=strict)
+
+            if missing:
+                logger.warning("Missing parameters when loading %s: %s", path, missing)
+            if unexpected:
+                logger.warning(
+                    "Unexpected parameters when loading %s: %s", path, unexpected
+                )
+
+            instance.to(device)
+            return instance
+
+        return super(MultiLabelClassificationNet, cls).load_from_checkpoint(
+            checkpoint_path, *args, **kwargs
+        )
+
+    def forward(self, x):  # pylint: disable=arguments-differ
+        x = self.model(x)
+        return x
+
+    def training_step(self, batch, _batch_idx):  # pylint: disable=arguments-differ
+        x, y = batch
+        y_pred = self(x)
+        loss = self.criterion(y_pred, y)
+        self.log(
+            "train/loss", loss, on_step=METRICS_ON_STEP, on_epoch=True, prog_bar=True
+        )
+
+        preds = np.array(self.sigm(y_pred).cpu() > 0.5, dtype=float)
+
+        return {"loss": loss, "preds": preds, "targets": y}
+
+    def validation_step(self, batch, _batch_idx):  # pylint: disable=arguments-differ
+        x, y = batch
+        y_pred = self(x)
+        loss = self.criterion(y_pred, y)
+        self.log("val/loss", loss, on_epoch=True, prog_bar=True)
+
+        preds = np.array(self.sigm(y_pred).cpu() > 0.5, dtype=float)
+        self.val_preds.append(preds)
+        self.val_targets.append(y.cpu().numpy())
+
+        return {"loss": loss, "preds": preds, "targets": y}
+
+    def validation_epoch_end(self, _outputs):
+        """Called at the end of validation to aggregate outputs"""
+        val_preds_np = np.concatenate(self.val_preds)
+        val_targets_np = np.concatenate(self.val_targets)
+
+        metrics = calculate_metrics(val_preds_np, val_targets_np, threshold=0.5)
+        for key, metric_value in metrics.items():
+            if isinstance(metric_value, np.ndarray):
+                processed_value = metric_value.tolist()
+            elif isinstance(metric_value, (list, tuple)):
+                processed_value = list(metric_value)
+            else:
+                processed_value = float(metric_value)
+
+            if isinstance(processed_value, list):
+                for i, single_value in enumerate(processed_value):
+                    name = "val/" + f"{key}/{self.labels[i]}"
+                    self.log(
+                        name,
+                        float(single_value),
+                        on_epoch=True,
+                        on_step=METRICS_ON_STEP,
+                        prog_bar=False,
+                    )
+            else:
+                name = "val/" + f"{key}"
+                self.log(
+                    name,
+                    float(processed_value),
+                    on_epoch=True,
+                    on_step=METRICS_ON_STEP,
+                    prog_bar=True,
+                )
+
+        self.val_preds = []
+        self.val_targets = []
+
+    def configure_optimizers(self):
+        """Choose what optimizers and learning-rate schedulers to use in your optimization.
+        Normally you'd need one. But in the case of GANs or similar you might have multiple.
+
+        See examples here:
+            https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers
+        """
+        optimizer = torch.optim.SGD(
+            self.parameters(), self.lr, momentum=0.5, weight_decay=self.weight_decay
+        )
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer,
+            T_0=20,
+        )
+
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": lr_scheduler,
+            "monitor": "val/loss",
+        }

endoreg_db/utils/ai/postprocess.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+
+def concat_pred_dicts(pred_dicts):
+    """Shoulkd be a list of dictionaries with the same keys"""
+    assert len(pred_dicts) > 0
+    keys = pred_dicts[0].keys()
+
+    merged_predictions = {key: [] for key in keys}
+    for p in pred_dicts:
+        for key in p.keys():
+            merged_predictions[key].append(p[key])
+
+    for key in merged_predictions.keys():
+        merged_predictions[key] = np.array(merged_predictions[key])
+
+    return merged_predictions
+
+
+def make_smooth_preds(prediction_array, window_size_s=1, fps=50):
+    window_size = int(window_size_s * fps)
+    smooth_prediction_array = np.convolve(
+        prediction_array, np.ones(window_size) / window_size, mode="valid"
+    )
+    return smooth_prediction_array
+
+
+def find_true_pred_sequences(predictions):
+    """
+    Efficiently finds sequences of 'outside' predictions in the binary predictions array using NumPy.
+
+    Args:
+    predictions (np.array): An array of boolean values, where True represents an 'outside' image
+                            and False represents an 'inside' image.
+
+    Returns:
+    list of tuples: A list where each tuple represents a sequence of 'outside' predictions,
+                    with the first element as the start index and the second element as the stop index.
+    """
+    # Identify where the value changes in the binary array (from False to True or True to False)
+    change_indices = np.where(np.diff(predictions.astype(int)) != 0)[0]
+
+    # Since diff reduces the length by 1, we adjust indices to align with the original array
+    change_indices += 1
+
+    # If the first element is 'outside', prepend a 0 to indicate the start
+    if predictions[0]:
+        change_indices = np.insert(change_indices, 0, 0)
+
+    # If the last element is 'outside', append the length of the array to indicate the end
+    if predictions[-1]:
+        change_indices = np.append(change_indices, predictions.size)
+
+    # Extract the 'outside' sequences by slicing the change_indices array in steps of two
+    outside_sequences = [
+        (change_indices[i], change_indices[i + 1] - 1)
+        for i in range(0, len(change_indices), 2)
+    ]
+
+    # make sure the result is serializable
+    outside_sequences = [(int(start), int(stop)) for start, stop in outside_sequences]
+
+    return outside_sequences

endoreg_db/utils/ai/predict.py

@@ -0,0 +1,293 @@
+"""
+This module contains the Classifier class for making predictions using a trained model.
+"""
+
+import json
+import torch
+from torch.utils.data import DataLoader
+from torch import nn
+import numpy as np
+from tqdm import tqdm
+from icecream import ic
+from .inference_dataset import InferenceDataset
+from .postprocess import concat_pred_dicts, make_smooth_preds, find_true_pred_sequences
+
+sample_config = {
+    # mean and std for normalization
+    "mean": (0.45211223, 0.27139644, 0.19264949),
+    "std": (0.31418097, 0.21088019, 0.16059452),
+    # Image Size
+    "size_x": 716,
+    "size_y": 716,
+    # how to wrangle axes of the image before putting them in the network
+    "axes": [2, 0, 1],  # 2,1,0 for opencv
+    "batchsize": 16,
+    "num_workers": 0,  # always 1 for Windows systems # FIXME: fix celery crash if multiprocessing
+    # maybe add sigmoid after prediction?
+    "activation": nn.Sigmoid(),
+    "labels": [
+        "appendix",
+        "blood",
+        "diverticule",
+        "grasper",
+        "ileocaecalvalve",
+        "ileum",
+        "low_quality",
+        "nbi",
+        "needle",
+        "outside",
+        "polyp",
+        "snare",
+        "water_jet",
+        "wound",
+    ],
+}
+
+
+class Classifier:
+    def __init__(self, model=None, config=None, verbose=False):
+        if config is None:
+            config = sample_config.copy()
+        self.config = config
+        self.model = model
+        self.verbose = verbose
+
+    def pipe(self, paths, crops, verbose=None):
+        """
+        Processes input data through the model pipeline and returns predictions.
+        Args:
+            paths (list): List of file paths to the input data.
+            crops (list): List of crop regions for the input data.
+            verbose (bool, optional): If True, prints detailed logs. Defaults to None.
+        Returns:
+            list: Predictions generated by the model.
+        """
+
+        if verbose is None:
+            verbose = self.verbose
+
+        dataset = InferenceDataset(paths, crops, self.config)
+        if verbose:
+            ic("Dataset created")
+            
+        use_cuda = torch.cuda.is_available()
+        
+        dl = DataLoader(
+            dataset=dataset,
+            batch_size=self.config["batchsize"],
+            num_workers=self.config["num_workers"],
+            shuffle=False,
+            pin_memory=use_cuda,
+        )
+        if verbose:
+            ic("Dataloader created")
+
+        predictions = []
+
+        with torch.inference_mode():
+            if self.verbose:
+                ic("Starting inference")
+
+            # Ensure model exists
+            if self.model is None:
+                raise ValueError("Model is not loaded")
+
+            # Use the device the model is currently on, with fallback to CPU
+            try:
+                # Check what device the model parameters are on
+                model_device = next(self.model.parameters()).device
+                device = model_device
+                if verbose:
+                    print(f"Using device: {device}")
+            except StopIteration:
+                # Model has no parameters, default to CPU
+                device = torch.device("cpu")
+                if verbose:
+                    print("Model has no parameters, defaulting to CPU")
+            except Exception as e:
+                # Any other issue, fall back to CPU
+                device = torch.device("cpu")
+                if verbose:
+                    print(f"Device detection failed, using CPU: {e}")
+
+            # Ensure model is in eval mode
+            self.model.eval()
+
+            for batch in tqdm(dl):
+                batch = batch.to(device, non_blocking=True)
+                prediction = self.model(batch)
+                prediction = (
+                    self.config["activation"](prediction).cpu().tolist()
+                )  # .numpy().tolist()
+                predictions += prediction
+
+        return predictions
+
+    def __call__(self, image, crop=None):
+        return self.pipe([image], [crop])
+
+    def readable(self, predictions):
+        """
+        Converts a list of predictions into a readable dictionary format.
+        Args:
+            predictions (list): A list of prediction values.
+        Returns:
+            dict: A dictionary where the keys are labels from the configuration
+                  and the values are the corresponding predictions. If a prediction
+                  value is of type numpy.int64, it is converted to a standard int.
+        """
+        readable_dict = {
+            label: prediction
+            for label, prediction in zip(self.config["labels"], predictions)
+        }
+
+        # when dumping to json we need to convert numpy.int64 to int
+        readable_dict = {
+            key: int(value) if isinstance(value, np.int64) else value
+            for key, value in readable_dict.items()
+        }
+
+        return readable_dict
+
+    def get_prediction_dict(self, predictions, paths):
+        """
+        Constructs a dictionary containing prediction results.
+        Args:
+            predictions (list): A list of prediction results.
+            paths (list): A list of paths corresponding to the predictions.
+        Returns:
+            dict: A dictionary with the following keys:
+                - "labels": The labels from the configuration.
+                - "paths": The provided paths.
+                - "predictions": The provided predictions.
+        """
+
+        json_dict = {
+            "labels": self.config["labels"],
+            "paths": paths,
+            "predictions": predictions,
+        }
+
+        return json_dict
+
+    def get_prediction_json(self, predictions, paths, json_target_path: str = None):
+        """
+        Saves predictions to a JSON file.
+
+        Args:
+            predictions (list): A list of prediction results.
+            paths (list): A list of paths corresponding to the predictions.
+            json_target_path (str, optional): The path to save the JSON file. Defaults to None.
+        """
+        if not json_target_path:
+            json_target_path = "predictions.json"
+
+        json_dict = self.get_prediction_dict(predictions, paths)
+
+        with open(json_target_path, "w", encoding="utf-8") as f:
+            json.dump(json_dict, f)
+
+        if self.verbose:
+            ic(f"Saved predictions to {json_target_path}")
+
+    def post_process_predictions(
+        self, pred_dicts, window_size_s=1, fps=50, min_seq_len_s=0.5
+    ):
+        """
+        pred_dicts: list of dictionaries with the same keys
+        window_size_s: size of the window in seconds for smoothing
+        fps: frames per second
+        min_seq_len_s: minimum length of a sequence in seconds
+
+        Returns:
+        predictions: concatenated predictions
+        smooth_predictions: smoothed predictions
+        binary_predictions: binary predictions
+        raw_sequences: raw sequences
+        filtered_sequences: filtered sequences
+        """
+        # Concatenate the predictions
+        predictions = concat_pred_dicts(pred_dicts)
+
+        smooth_predictions = {key: [] for key in predictions.keys()}
+        for key in predictions.keys():
+            smooth_predictions[key] = make_smooth_preds(
+                predictions[key], window_size_s=window_size_s, fps=fps
+            )
+
+        binary_predictions = {}
+        for key in smooth_predictions.keys():
+            binary_predictions[key] = np.array(
+                [p > 0.5 for p in smooth_predictions[key]]
+            )
+
+        raw_sequences = {}
+        for key, value in binary_predictions.items():
+            raw_sequences[key] = find_true_pred_sequences(value)
+
+        filtered_sequences = {}
+        min_seq_len = int(min_seq_len_s * fps)
+        for key, sequences in raw_sequences.items():
+            filtered_sequences[key] = [
+                s for s in sequences if s[1] - s[0] > min_seq_len
+            ]
+
+        return (
+            predictions,
+            smooth_predictions,
+            binary_predictions,
+            raw_sequences,
+            filtered_sequences,
+        )
+
+    def post_process_predictions_serializable(
+        self, pred_dicts, window_size_s=1, fps=50, min_seq_len_s=0.5
+    ):
+        """
+        Post-processes prediction dictionaries to make them serializable.
+        This method takes prediction dictionaries, processes them to ensure all
+        elements are serializable (e.g., converting numpy arrays to lists), and
+        organizes the results into a dictionary of dictionaries.
+        Args:
+            pred_dicts (list): List of prediction dictionaries.
+            window_size_s (int, optional): Window size in seconds. Defaults to 1.
+            fps (int, optional): Frames per second. Defaults to 50.
+            min_seq_len_s (float, optional): Minimum sequence length in seconds. Defaults to 0.5.
+        Returns:
+            dict: A dictionary containing processed prediction results with the following keys:
+                - "predictions": Processed predictions.
+                - "smooth_predictions": Smoothed predictions.
+                - "binary_predictions": Binary predictions.
+                - "raw_sequences": Raw sequences.
+                - "filtered_sequences": Filtered sequences.
+        """
+
+        result = self.post_process_predictions(
+            pred_dicts, window_size_s, fps, min_seq_len_s
+        )
+
+        for i, _dict in enumerate(result):
+            _keys = list(_dict.keys())
+            for key in _keys:
+                # if numpy array
+                if hasattr(_dict[key], "tolist"):
+                    result[i][key] = _dict[key].tolist()
+
+                # check if list of tuples
+                # if so, make sure each tuple has 2 elements and split to two lists (start, stop)
+                if all(isinstance(x, tuple) for x in _dict[key]):
+                    if all(len(x) == 2 for x in _dict[key]):
+                        result[i][f"{key}_start"] = [int(x[0]) for x in _dict[key]]
+                        result[i][f"{key}_stop"] = [int(x[1]) for x in _dict[key]]
+                        del result[i][key]
+
+        # make dict of dicts
+        result_dict = {
+            "predictions": result[0],
+            "smooth_predictions": result[1],
+            "binary_predictions": result[2],
+            "raw_sequences": result[3],
+            "filtered_sequences": result[4],
+        }
+
+        return result_dict

endoreg_db/utils/ai/preprocess.py

@@ -0,0 +1,76 @@
+import numpy as np
+from PIL import Image, ImageOps  # Import the required modules from Pillow
+
+
+def crop_img(img, crop):
+    """
+    Crops the image based on the specified dimensions and adds padding to maintain aspect ratio.
+
+    Parameters:
+        img: PIL Image object.
+        crop: Tuple of (ymin, ymax, xmin, xmax) specifying the crop area.
+
+    Returns:
+        PIL Image object that has been cropped and padded as necessary.
+    """
+    # Convert crop dimensions to Pillow format: left, upper, right, lower
+    ymin, ymax, xmin, xmax = crop
+    img_cropped = img.crop((xmin, ymin, xmax, ymax))
+
+    # Calculate the new size and the required padding
+    width, height = img_cropped.size
+    delta = width - height
+
+    if delta > 0:
+        padding = (
+            0,
+            abs(delta) // 2,
+            0,
+            abs(delta) - abs(delta) // 2,
+        )  # (left, top, right, bottom)
+    elif delta < 0:
+        padding = (abs(delta) // 2, 0, abs(delta) - abs(delta) // 2, 0)
+    else:
+        padding = (0, 0, 0, 0)
+
+    # Pad the image to make it square
+    img_padded = ImageOps.expand(img_cropped, padding)
+
+    return img_padded
+
+
+class Cropper:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self, img, crop=None, scale=None, scale_method=Image.Resampling.LANCZOS
+    ):
+        """
+        Applies cropping and scaling transformations to the input image.
+
+        Parameters:
+            img: PIL Image object or numpy array of the image.
+            crop: Optional tuple specifying the cropping area (y_min, y_max, x_min, x_max).
+            scale: Optional tuple specifying the new size (width, height).
+            scale_method: Resampling method used for scaling (default is Image.Resampling.LANCZOS).
+
+        Returns:
+            Numpy array of the processed image.
+        """
+        # Convert numpy array to PIL Image if necessary
+        if isinstance(img, np.ndarray):
+            img = Image.fromarray(img.astype("uint8"), "RGB")
+
+        if crop is not None:
+            img = crop_img(img, crop)
+        else:
+            raise Exception("Automatic crop detection not implemented yet")
+
+        if scale is not None:
+            img = img.resize(scale, resample=scale_method)
+
+        # Convert PIL Image back to numpy array
+        img = np.array(img)
+
+        return img