dataset_ref.py

import random

import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
import torch.utils.data as data
from PIL import Image
import os
import os.path
from typing import Callable, cast, Dict, List, Optional, Tuple, Union
from pathlib import Path
from collections import OrderedDict

IMG_EXTENSIONS = ".npy"

class Threshold_noise:
    """Remove intensities under given threshold"""

    def __init__(self, threshold=100.):
        self.threshold = threshold

    def __call__(self, x):
        return torch.where((x <= self.threshold), 0.,x)

class Log_normalisation:
    """Log normalisation of intensities"""

    def __init__(self, eps=1e-5):
        self.epsilon = eps

    def __call__(self, x):
        return torch.log(x+1+self.epsilon)/torch.log(torch.max(x)+1+self.epsilon)

class Random_shift_rt():
    pass


def default_loader(path):
    return Image.open(path).convert('RGB')

def npy_loader(path):
    sample = torch.from_numpy(np.load(path))
    sample = sample.unsqueeze(0)
    return sample

def remove_aer_ana(l):
    l = map(lambda x : x.split('_')[0],l)
    return list(OrderedDict.fromkeys(l))

def make_dataset_custom(
    directory: Union[str, Path],
    class_to_idx: Optional[Dict[str, int]] = None,
    extensions: Optional[Union[str, Tuple[str, ...]]] = IMG_EXTENSIONS,
    is_valid_file: Optional[Callable[[str], bool]] = None,
    allow_empty: bool = False,
) -> List[Tuple[str, str, int]]:
    """Generates a list of samples of a form (path_to_sample, class).

    See :class:`DatasetFolder` for details.

    Note: The class_to_idx parameter is here optional and will use the logic of the ``find_classes`` function
    by default.
    """
    directory = os.path.expanduser(directory)

    if class_to_idx is None:
        _, class_to_idx = torchvision.datasets.folder.find_classes(directory)
    elif not class_to_idx:
        raise ValueError("'class_to_index' must have at least one entry to collect any samples.")

    both_none = extensions is None and is_valid_file is None
    both_something = extensions is not None and is_valid_file is not None
    if both_none or both_something:
        raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time")

    if extensions is not None:

        def is_valid_file(x: str) -> bool:
            return torchvision.datasets.folder.has_file_allowed_extension(x, extensions)  # type: ignore[arg-type]

    is_valid_file = cast(Callable[[str], bool], is_valid_file)

    instances = []
    available_classes = set()
    for target_class in sorted(class_to_idx.keys()):
        class_index = class_to_idx[target_class]
        target_dir = os.path.join(directory, target_class)
        if not os.path.isdir(target_dir):
            continue
        for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
            fnames_base = remove_aer_ana(fnames)
            for fname in sorted(fnames_base):
                fname_ana = fname+'_ANA.npy'
                fname_aer = fname + '_AER.npy'
                path_ana = os.path.join(root, fname_ana)
                path_aer = os.path.join(root, fname_aer)
                if is_valid_file(path_ana) and is_valid_file(path_aer) and os.path.isfile(path_ana) and os.path.isfile(path_aer):
                    item = path_aer, path_ana, class_index
                    instances.append(item)

                    if target_class not in available_classes:
                        available_classes.add(target_class)

    empty_classes = set(class_to_idx.keys()) - available_classes
    if empty_classes and not allow_empty:
        msg = f"Found no valid file for the classes {', '.join(sorted(empty_classes))}. "
        if extensions is not None:
            msg += f"Supported extensions are: {extensions if isinstance(extensions, str) else ', '.join(extensions)}"
        raise FileNotFoundError(msg)

    return instances


class ImageFolderDuo(data.Dataset):
    def __init__(self, root, transform=None, target_transform=None,
                 flist_reader=make_dataset_custom, loader=npy_loader, ref_dir = None, positive_prop=None):
        self.root = root
        self.imlist = flist_reader(root)
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader
        self.classes = torchvision.datasets.folder.find_classes(root)[0]
        self.ref_dir = ref_dir
        self.positive_prop = positive_prop

    def __getitem__(self, index):
        impathAER, impathANA, target = self.imlist[index]
        imgAER = self.loader(impathAER)
        imgANA = self.loader(impathANA)
        if self.positive_prop is not None:
            i = random.randint(0,100)
            if i < self.positive_prop:
                label_ref = target
            else :
                label_ref = np.random.randint(0,len(self.classes)-1) #can be postive too (border effect)
        else :
            label_ref = np.random.randint(0,len(self.classes)-1)
        class_ref = self.classes[label_ref]
        path_ref = self.ref_dir +'/'+ class_ref + '.npy'
        img_ref = self.loader(path_ref)
        if self.transform is not None:
            imgAER = self.transform(imgAER)
            imgANA = self.transform(imgANA)
            img_ref = self.transform(img_ref)
        target = 0 if target == label_ref else 1
        return imgAER, imgANA, img_ref, target

    def __len__(self):
        return len(self.imlist)

def load_data_duo(base_dir_train, base_dir_test, batch_size, shuffle=True, noise_threshold=0, ref_dir = None, positive_prop=None):

    train_transform = transforms.Compose(
        [transforms.Resize((224, 224)),
         Threshold_noise(noise_threshold),
         Log_normalisation(),
         transforms.Normalize(0.5, 0.5)])
    print('Default train transform')

    val_transform = transforms.Compose(
        [transforms.Resize((224, 224)),
         Threshold_noise(noise_threshold),
         Log_normalisation(),
         transforms.Normalize(0.5, 0.5)])
    print('Default val transform')

    train_dataset = ImageFolderDuo(root=base_dir_train, transform=train_transform, ref_dir = ref_dir, positive_prop=positive_prop)
    val_dataset = ImageFolderDuo_Batched(root=base_dir_test, transform=val_transform, ref_dir = ref_dir)

    data_loader_train = data.DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=0,
        collate_fn=None,
        pin_memory=False,
    )

    data_loader_test = data.DataLoader(
        dataset=val_dataset,
        batch_size=1,
        shuffle=shuffle,
        num_workers=0,
        collate_fn=None,
        pin_memory=False,
    )

    return data_loader_train, data_loader_test


class ImageFolderDuo_Batched(data.Dataset):
    def __init__(self, root, transform=None, target_transform=None,
                 flist_reader=make_dataset_custom, loader=npy_loader, ref_dir = None):
        self.root = root
        self.imlist = flist_reader(root)
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader
        self.classes = torchvision.datasets.folder.find_classes(root)[0]
        self.ref_dir = ref_dir

    def __getitem__(self, index):
        impathAER, impathANA, target = self.imlist[index]
        imgAER = self.loader(impathAER)
        imgANA = self.loader(impathANA)
        img_refs = []
        label_refs = []
        for ind_ref in range(len(self.classes)):
            class_ref = self.classes[ind_ref]
            target_ref = 0 if target == ind_ref else 1
            path_ref = self.ref_dir +'/'+ class_ref + '.npy'
            img_ref = self.loader(path_ref)
            if self.transform is not None:
                img_ref = self.transform(img_ref)
            img_refs.append(img_ref)
            label_refs.append(target_ref)
        if self.transform is not None:
            imgAER = self.transform(imgAER)
            imgANA = self.transform(imgANA)

        batched_im_ref = torch.concat(img_refs,dim=0)
        batched_label = torch.tensor(label_refs)
        batched_imgAER = imgAER.repeat(len(self.classes),1,1)
        batched_imgANA = imgANA.repeat(len(self.classes),1,1)

        return batched_imgAER, batched_imgANA, batched_im_ref, batched_label

    def __len__(self):
        return len(self.imlist)