endoreg-db 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

endoreg_db/models/ai_model/lightning/inference_dataset.py

@@ -0,0 +1,53 @@
+from torch.utils.data import Dataset
+import torch
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+from .preprocess import Cropper
+
+class InferenceDataset(Dataset):
+    def __init__(self, paths, crops, config):
+        self.paths = paths
+        self.crops = crops
+        self.cropper = Cropper()  # Assuming Cropper can work with NumPy arrays
+        self.config = config
+
+        # Initialize the image transformations using torchvision
+        self.transforms = transforms.Compose([
+            # Convert PIL image to PyTorch tensor
+            transforms.ToTensor(),
+            # Normalize the image using the provided mean and std
+            transforms.Normalize(mean=self.config["mean"], std=self.config["std"])
+        ])
+        
+    def __len__(self):
+        # Returns the total number of samples
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        # Open the image with Pillow
+        with Image.open(self.paths[idx]) as pil_image:
+            # Convert the image to RGB to ensure 3 channels
+            pil_image = pil_image.convert('RGB')
+
+        # Get the corresponding crop for the current image
+        crop = self.crops[idx]
+
+        # Crop the image based on the provided crop parameters and convert to numpy for cropping
+        cropped = self.cropper(
+            np.array(pil_image),  # Convert PIL image to numpy array for cropping
+            crop,
+            scale=[
+                self.config["size_x"],
+                self.config["size_y"]
+            ]    
+        )
+
+        # Convert cropped numpy array back to PIL image for torchvision transforms
+        cropped_pil = Image.fromarray(cropped.astype('uint8'), 'RGB')
+        
+        # Apply the transformations
+        img = self.transforms(cropped_pil)
+
+
+        return img

endoreg_db/models/ai_model/lightning/multilabel_classification_net.py

@@ -0,0 +1,155 @@
+import torch
+from torchvision import models
+import torch.nn as nn
+from pytorch_lightning import LightningModule
+import numpy as np
+from sklearn.metrics import precision_score, recall_score, f1_score
+
+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
+        ),
+    }
+
+
+class MultiLabelClassificationNet(LightningModule):
+    def __init__(
+        self,
+        labels=["ASD"],
+        lr=6e-3,
+        weight_decay=0.001,
+        pos_weight=2,
+        model_type="EfficientNetB4",
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        self.model_type = "RegNetX800MF"  # model_type
+        self.labels = labels
+        self.n_classes = len(labels)
+        self.val_preds = []
+        self.val_targets = []
+        self.pos_weight = pos_weight
+        self.weight_decay = weight_decay
+        self.lr = lr
+        self.sigm = nn.Sigmoid()
+
+        if model_type == "EfficientNetB4":
+            self.model = models.efficientnet_b4(pretrained=True)
+            num_ftrs = self.model.classifier[1].in_features
+            self.model.classifier[1] = nn.Linear(num_ftrs, len(labels))
+
+        elif model_type == "RegNetX800MF":
+            self.model = models.regnet_x_800mf(pretrained=True)
+            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))
+        )
+
+    def forward(self, x):
+        """Forward pass"""
+        x = self.model(x)
+        return x
+
+    def training_step(self, batch, batch_idx):
+        """Training step"""
+        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):
+        """Validation step"""
+        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):
+        """Validation epoch end"""
+        self.val_preds = np.concatenate([_ for _ in self.val_preds])
+        self.val_targets = np.concatenate([_ for _ in self.val_targets])
+
+        metrics = calculate_metrics(self.val_preds, self.val_targets, threshold=0.5)
+        for key, value in metrics.items():
+            value = value.tolist()
+            if isinstance(value, list):
+                for i, _value in enumerate(value):
+                    name = "val/" + f"{key}/{self.labels[i]}"
+                    self.log(
+                        name,
+                        _value,
+                        on_epoch=True,
+                        on_step=METRICS_ON_STEP,
+                        prog_bar=False,
+                    )
+            else:
+                name = "val/" + f"{key}"
+                self.log(
+                    name, 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, verbose=True
+        )
+
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": lr_scheduler,
+            "monitor": "val/loss",
+        }

endoreg_db/models/ai_model/lightning/postprocess.py

@@ -0,0 +1,53 @@
+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)]
+
+    return outside_sequences
+

endoreg_db/models/ai_model/lightning/predict.py

@@ -0,0 +1,172 @@
+from .inference_dataset import InferenceDataset
+from torch.utils.data import DataLoader
+import torch
+from torch import nn
+import json
+from .postprocess import (
+    concat_pred_dicts,
+    make_smooth_preds,
+    find_true_pred_sequences
+)
+import numpy as np
+
+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=sample_config, verbose = False):
+        self.config = config
+        self.model = model
+        self.verbose = verbose
+
+    def pipe(self, paths, crops, verbose = None):
+        if verbose is None:
+            verbose = self.verbose
+
+        dataset = InferenceDataset(paths, crops, self.config)
+        if verbose:
+            print("Dataset created")
+
+        dl = DataLoader(
+            dataset=dataset,
+            batch_size=self.config["batchsize"],
+            num_workers=self.config["num_workers"],
+            shuffle = False,
+            pin_memory=True
+        )
+        if verbose:
+            print("Dataloader created")
+
+        predictions = []
+
+        with torch.inference_mode():
+            if self.verbose:
+                print("Starting inference")
+            for i,batch in enumerate(dl):
+                prediction = self.model(batch.cuda())
+                prediction = self.config["activation"](prediction).cpu().tolist()#.numpy().tolist()
+                predictions += prediction
+                if self.verbose and i==0:
+                    print("First batch done")
+
+        return predictions
+
+    def __call__(self, image, crop=None):
+        return self.pipe([image], [crop])
+
+    def readable(self, predictions):
+        return {label: prediction for label, prediction in zip(self.config["labels"], predictions)}
+    
+    def get_prediction_dict(self, predictions, paths):
+        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):
+        if not json_target_path:
+            json_target_path = "predictions.json"
+
+        json_dict = self.get_prediction_dict(predictions, paths)
+
+        with open(json_target_path, 'w') as f:
+            json.dump(json_dict, f)
+        
+        if self.verbose:
+            print(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 in binary_predictions.keys():
+            raw_sequences[key] = find_true_pred_sequences(binary_predictions[key])
+
+        filtered_sequences = {}
+        min_seq_len = int(min_seq_len_s * fps)
+        for key in raw_sequences.keys():
+            filtered_sequences[key] = [s for s in raw_sequences[key] 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
+    ):
+        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/models/ai_model/lightning/prediction_visualizer.py

@@ -0,0 +1,55 @@
+import streamlit as st
+import matplotlib.pyplot as plt
+import json
+from PIL import Image
+import argparse
+
+
+# Load the JSON file
+@st.cache_data
+def load_data(file_path):
+    with open(file_path, 'r') as f:
+        data = json.load(f)
+    return data
+
+# Streamlit App
+def main(file_path):
+    st.title("Visualize Predictions")
+
+    # Load the data
+    data = load_data(file_path)
+
+    # Dropdown to select label
+    selected_label = st.selectbox('Select Label', options=data['labels'])
+
+    # Get the index of the selected label
+    label_index = data['labels'].index(selected_label)
+
+    # Extract predictions for the selected label
+    selected_predictions = [pred[label_index] for pred in data['predictions']]
+
+    # Line plot for the selected label's prediction values
+    st.subheader(f"Line plot for {selected_label}")
+
+    fig, ax = plt.subplots(figsize=(10, 6))  # Explicitly create a figure and axis
+    ax.plot(selected_predictions)
+    ax.set_xlabel('Frame')
+    ax.set_ylabel('Prediction Value')
+    ax.set_title(f'Predictions for {selected_label}')
+
+    st.pyplot(fig)  # Pass the figure object explicitly
+
+
+    # Slider to select frame
+    frame_idx = st.slider('Select a frame', 0, len(data['paths']) - 1)
+
+    # Display the frame
+    image_path = data['paths'][frame_idx]
+    image = Image.open(image_path)
+    st.image(image, caption=f"Frame {frame_idx}", use_column_width=True)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="Visualize Predictions")
+    parser.add_argument("--file", type=str, help="Path to JSON file containing predictions", required=True)
+    args = parser.parse_args()
+    main(args.file)

endoreg_db/models/ai_model/lightning/preprocess.py

@@ -0,0 +1,68 @@
+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

endoreg_db/models/ai_model/lightning/run_visualizer.py

@@ -0,0 +1,21 @@
+# agl_visualization/run_visualizer.py
+
+import subprocess
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser(description="Visualize Predictions")
+    parser.add_argument("--file", type=str, help="Path to JSON file containing predictions", required=True)
+    args = parser.parse_args()
+
+    # Replace this with the path where prediction_visualizer.py resides
+    streamlit_script_path = "agl_predict_endo_frame/prediction_visualizer.py"
+    
+    # Using subprocess to run the Streamlit command
+    subprocess.run(f'streamlit run {streamlit_script_path} -- --file {args.file}', shell=True)
+
+# if __name__ == "__main__":
+    # parser = argparse.ArgumentParser(description="Visualize Predictions")
+    # parser.add_argument("--file", type=str, help="Path to JSON file containing predictions", required=True)
+    # args = parser.parse_args()
+    # main(args.file)