import os
import torch
from torchvision.datasets.folder import make_dataset
from pathlib import Path
from torch.utils.data import Dataset, DataLoader, Subset
from lightning.pytorch import LightningDataModule
import numpy as np
import cv2
from typing import Optional, Sequence, Tuple, Union, List, Any, Callable
from sklearn.model_selection import KFold, StratifiedKFold
import torchvision
import math
import defs
import csv
import pandas as pd
from enum import Enum
import shutil
from ultralytics import YOLO
import itertools
def _find_classes(dir):
classes = [d.name for d in os.scandir(dir) if d.is_dir()]
classes.sort()
class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
return classes, class_to_idx
def get_samples(root, extensions=(".mp4", ".avi")):
_, class_to_idx = _find_classes(root)
return make_dataset(root, class_to_idx, extensions=extensions)
class RARP_DatasetCreator():
def _removeBlackBorder(self, image):
image = np.array(image)
copyImg = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2HSV)
h = copyImg[:,:,0]
mask = np.ones(h.shape, dtype=np.uint8) * 255
th = (25, 175)
mask[(h > th[0]) & (h < th[1])] = 0
copyImg = cv2.cvtColor(copyImg, cv2.COLOR_HSV2BGR)
resROI = cv2.bitwise_and(copyImg, copyImg, mask=mask)
image_gray = cv2.cvtColor(resROI, cv2.COLOR_BGR2GRAY)
_, thresh = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
morph = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
contours = cv2.findContours(morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours = contours[0] if len(contours) == 2 else contours[1]
bigCont = max(contours, key=cv2.contourArea)
x, y, w, h = cv2.boundingRect(bigCont)
crop = image[y : y + h, x : x + w]
return crop
def _crop(self, image:np, x:tuple, y:tuple):
return image[y[0]:y[1], x[0]:x[1], :]
def _sliceImage(self, img:np, size:float):
h, w, _ = img.shape
xw = math.trunc(w*size)
xh = math.trunc(h*size)
ul = self._crop(img, (0, xw), (0, xh))
ur = self._crop(img, (w-xw, w), (0, xh))
dl = self._crop(img, (0, xw), (h-xh, h))
dr = self._crop(img, (w-xw, w), (h-xh, h))
return [ul, ur, dl, dr]
def _split(self, a, n):
k, m = divmod(len(a), n)
return (a[i*k+min(i, m):(i+1)*k+min(i+1, m)] for i in range(n))
def ROI_Extract_YOLO(self, YoloModel:YOLO, image, threshold=0.75):
listROIs = []
res = YoloModel(image, stream=True)
for r in res:
pred = r.boxes.conf.cpu().numpy()
for i, box in enumerate(r.boxes.xyxy.cpu().numpy()):
if pred[i] >= threshold:
box = box.astype(int)
x, y = (box[0], box[1])
xw, yh = (box[2], box[3])
listROIs.append(image[y:yh, x:xw])
return listROIs
def __init__(self,
RootPath = "",
extension="tiff",
FoldSeed=505,
createFile=False,
SavePath="",
Fold:int=None,
preResize=None,
RGBGama=False,
removeBlackBar=True,
SegImage:float=None,
Num_Img_Slices:int = 4,
SegmentClass:int = None,
colorSpace:int=None,
ROI_Yolo:YOLO=None,
thresholdYolo_Accuracy:float=0.75) -> None:
root = Path(RootPath)
lMean = []
lStd = []
NO_NVB = []
NVB = []
if createFile:
if len(SavePath) == 0:
raise Exception("If createFile is True, SavePath must have a value")
sPath = Path(SavePath + f"_seed_{FoldSeed}")
dumpImgs = sPath/"dump"
dumpImgs.mkdir(parents=True, exist_ok=True)
self.CVS_File = dumpImgs/"dataset.csv"
if not self.CVS_File.exists():
with open(self.CVS_File, "x", newline='') as csvfile:
writerOBJ = csv.writer(csvfile)
writerOBJ.writerow(["id", "label", "class", "path", "name", "mean_1", "mean_2", "mean_3", "std_1", "std_2", "std_3"])
id = 0
for i, file in enumerate(root.glob(f"**/*.{extension}")):
tempImg = cv2.imread(str(file), cv2.IMREAD_COLOR)
if removeBlackBar and ROI_Yolo is None:
tempImg = self._removeBlackBorder(tempImg)
if RGBGama:
tempImg = cv2.cvtColor(tempImg, cv2.COLOR_BGR2RGB)
if colorSpace is not None:
tempImg = cv2.cvtColor(tempImg, colorSpace)
if preResize is not None:
tempImg = cv2.resize(tempImg, preResize, interpolation=cv2.INTER_AREA)#Cambio de resolucion
tempImgList = [tempImg]
#ROI w/ YOLOv8
if ROI_Yolo is not None:
tempImgList = self.ROI_Extract_YOLO(ROI_Yolo, tempImg, thresholdYolo_Accuracy)
for k, tempImg in enumerate(tempImgList):
lineaCSV = [id, 0 if file.parent.name == "NOT_NVB" else 1, file.parent.name, (dumpImgs/f"Img{k}-{i}.npy").absolute(), file.name]
lMean.append(np.mean(tempImg, axis=tuple(range(tempImg.ndim-1))))
lStd.append(np.std(tempImg, axis=tuple(range(tempImg.ndim-1))))
lineaCSV += np.mean(tempImg, axis=tuple(range(tempImg.ndim-1))).tolist()
lineaCSV += np.std(tempImg, axis=tuple(range(tempImg.ndim-1))).tolist()
writerOBJ.writerow(lineaCSV)
#lista.append (np.mean(tempImg, axis=tuple(range(tempImg.ndim-1))))
np.save(dumpImgs/f"Img{k}-{i}.npy", tempImg)
if lineaCSV[1] == 0:
NO_NVB.append(id)
else:
NVB.append(id)
id += 1
if SegImage is not None:
if (SegmentClass is None) or (lineaCSV[1] == SegmentClass):
for j, newImg in enumerate(self._sliceImage(tempImg, SegImage)):
if j == Num_Img_Slices:
break
imgPath = dumpImgs/f"Img{k}-{i}-{j}.npy"
np.save(imgPath, newImg)
lineaCSV = [id, 0 if file.parent.name == "NOT_NVB" else 1, file.parent.name, (dumpImgs/f"Img{k}-{i}-{j}.npy").absolute(), file.name]
lineaCSV += np.mean(newImg, axis=tuple(range(newImg.ndim-1))).tolist()
lineaCSV += np.std(newImg, axis=tuple(range(newImg.ndim-1))).tolist()
writerOBJ.writerow(lineaCSV)
if lineaCSV[1] == 0:
NO_NVB.append(id)
else:
NVB.append(id)
id += 1
csvfile.close()
temp = np.asarray(lMean)
self.mean = list(np.mean(temp, axis=tuple(range(temp.ndim-1))))
temp = np.asarray(lStd)
self.std = list(np.mean(temp, axis=tuple(range(temp.ndim-1))))
else:
data = pd.read_csv(self.CVS_File)
self.mean = data[["mean_1", "mean_2", "mean_3"]].mean().to_list()
self.std = data[["std_1", "std_2", "std_3"]].mean().to_list()
NO_NVB = data.loc[data["label"] == 0]["id"].to_list()
NVB = data.loc[data["label"] == 1]["id"].to_list()
if Fold is not None:
self.Splits = []
splitsToSave = []
tempFoldsOrder = list(range(Fold))
for _ in range(Fold):
self.Splits.append(tempFoldsOrder)
tempFoldsOrder = tempFoldsOrder[1:] + tempFoldsOrder[:1]
if FoldSeed is not None:
np.random.seed(FoldSeed)
np.random.shuffle(NO_NVB)
np.random.shuffle(NVB)
NO_NVB_Folds = list (self._split(NO_NVB, Fold))
NVB_Fols = list(self._split(NVB, Fold))
setst = [
math.trunc(0.60 * Fold), math.trunc(0.20 * Fold), math.trunc(0.20 * Fold)
]
for data in self.Splits:
ultimo = 0
for s in setst:
tempArry = []
for fold in data[ultimo: ultimo + s]:
tempArry += NO_NVB_Folds[fold] + NVB_Fols[fold]
splitsToSave.append(tempArry)
ultimo += s
self.Folds_File = dumpImgs/"Folds.npy"
np.save(self.Folds_File, np.asarray(splitsToSave, dtype=object))
def CreateClases(self):
if not (self.CVS_File.parent/"FOLDS_CREATED").exists():
root = self.CVS_File.parent.parent
database = pd.read_csv(self.CVS_File, usecols=["id", "label", "class", "path"])
for _, row in database.iterrows():
pathOriginal = Path(row["path"])
pathNuevo = root/"dataset"/("NO_NVB" if row["label"] == 0 else "NVB")/pathOriginal.name
pathNuevo.parent.mkdir(parents=True, exist_ok=True)
shutil.copy(pathOriginal, pathNuevo)
with open(self.CVS_File.parent/"FOLDS_CREATED", "x") as file:
file.close()
for f in self.CVS_File.parent.glob("*.npy"):
f.unlink()
def ExtractROI_YOLO(self, YOLOModel:YOLO, thresholdYolo:float = 0.70):
if (self.CVS_File.parent/"FOLDS_CREATED").exists() and not (self.CVS_File.parent/"YOLO_ROI_CREATED").exists():
root = self.CVS_File.parent.parent
for img in root.glob(f"**/*.npy"):
if img.name == "Folds.npy":
continue
tempImgList = self.ROI_Extract_YOLO(YOLOModel, np.load(img), thresholdYolo)
for k, tempImg in enumerate(tempImgList):
if k == 0:
np.save(img, tempImg)
#else:
# newPath = img.parent/(img.name.replace(".npy", "") +f"-{k}.npy")
# np.save(newPath, tempImg)
with open(self.CVS_File.parent/"YOLO_ROI_CREATED", "x") as file:
file.close()
def CreateFolds(self):
if not (self.CVS_File.parent/"FOLDS_CREATED").exists():
root = self.CVS_File.parent.parent
database = pd.read_csv(self.CVS_File, usecols=["id", "label", "class", "path"])
arrFolds = np.load(self.Folds_File, allow_pickle=True)
for foldNum, splits in enumerate(np.array_split(arrFolds, len(arrFolds)/3)):
foldPath = root/f"fold_{foldNum}"
for datasetType, subSet in enumerate(splits):
SubfoldPath = foldPath/f"{RARP_DataSetType(datasetType).name}"
for _, row in database.loc[database["id"].isin(subSet)].iterrows():
PathOri = Path(row["path"])
PathImg = SubfoldPath/("NO_NVB" if row["label"] == 0 else "NVB")/PathOri.name
PathImg.parent.mkdir(parents=True, exist_ok=True)
shutil.copy(PathOri, PathImg)
with open(self.CVS_File.parent/"FOLDS_CREATED", "x") as file:
file.close()
for f in self.CVS_File.parent.glob("*.npy"):
f.unlink()
class RARP_DatasetFolder_ExtraData(torchvision.datasets.DatasetFolder):
def __init__(self,
root: str,
loader: Callable[[str], Any],
Extra_Data: pd.DataFrame,
Extra_Data_leg:int = 4,
extensions: Tuple[str, ...] | None = None,
transform: Callable[..., Any] | None = None,
target_transform: Callable[..., Any] | None = None,
is_valid_file: Callable[[str], bool] | None = None
) -> None:
super().__init__(root, loader, extensions, transform, target_transform, is_valid_file)
self.Extra_Data = Extra_Data
self.Extra_Data_leg = Extra_Data_leg
def __getitem__(self, index: int) -> Tuple[Any, Any]:
path, target = self.samples[index]
name = Path(path).name
Extra_data = self.Extra_Data[self.Extra_Data["raw_name"] == name].values.flatten().tolist()[:self.Extra_Data_leg]
Extra_data = torch.tensor(Extra_data)
#Loadres is 2 values
sample = self.loader(path)
if self.transform is not None:
sample = self.transform(sample)
if self.target_transform is not None:
target = self.target_transform(target)
return (sample, Extra_data), target
class RARP_DatasetFolder_DobleTransform(torchvision.datasets.DatasetFolder):
def __init__(self,
root: str,
loader: Callable[[str], Any],
extensions: Tuple[str, ...] | None = None,
transformT1: Callable[..., Any] | None = None,
transformT2: Callable[..., Any] | None = None,
target_transform: Callable[..., Any] | None = None,
is_valid_file: Callable[[str], bool] | None = None,
passOriginal: Callable[..., Any] | None = None,
) -> None:
super().__init__(root, loader, extensions, None, target_transform, is_valid_file)
self.T1 = transformT1
self.T2 = transformT2 if transformT2 is not None else transformT1
self.PassOrigianlImage = passOriginal
def __getitem__(self, index: int) -> Tuple[Any, Any]:
path, target = self.samples[index]
sample = self.loader(path)
if self.T1 is not None:
sample1 = self.T1(sample)
if self.T2 is not None:
sample2 = self.T2(sample)
if self.PassOrigianlImage is not None:
return (sample1, sample2, self.PassOrigianlImage(sample)), target
else:
return (sample1, sample2), target
class RARP_PreprocessCreator():
def removeBlackBorder(self, image):
image = np.array(image)
image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
_, thresh = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
morph = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
contours = cv2.findContours(morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours = contours[0] if len(contours) == 2 else contours[1]
bigCont = max(contours, key=cv2.contourArea)
x, y, w, h = cv2.boundingRect(bigCont)
crop = image[y : y + h, x : x + w]
return crop
def _crop(self, image:np, x:tuple, y:tuple):
return image[y[0]:y[1], x[0]:x[1], :]
def _sliceImage(self, img:np, size:float):
h, w, _ = img.shape
ul = self._crop(img, (0, math.trunc(w*size)), (0, math.trunc(h*size)))
ur = self._crop(img, (math.trunc(w*size), w), (0, math.trunc(h*size)))
dl = self._crop(img, (0, math.trunc(w*size)), (math.trunc(h*size), h))
dr = self._crop(img, (math.trunc(w*size), w), (math.trunc(h*size), h))
return [ul, ur, dl, dr]
def __init__(self,
RootPath = "",
extension="tiff",
FoldSeed=505,
createFile=False,
SavePath="",
Fold=False,
preResize=None,
RGBGama=False,
removeBlackBar=True,
SegImage:float=None) -> None:
super().__init__()
self.root = Path(RootPath)
self.mean = 0.0
self.std = 0.0
self.test_DataSet = []
self.val_DataSet = []
self.train_DataSet = []
self.nNVB = []
self.yNVB = []
lista = []
dist = {
'test': [5, 2], #[4, 3], [10, 4],
'val': [8, 3], #[4, 3], [10, 4],
'train': [28, 10] #[8, 6], [21, 7]
}
sPath = Path(SavePath + f"_seed_{FoldSeed}")
dumpImgs = sPath/"dump"
dumpImgs.mkdir(parents=True, exist_ok=True)
for i, file in enumerate(self.root.glob(f"**/*.{extension}")):
tempImg = cv2.imread(str(file), cv2.IMREAD_COLOR)
if removeBlackBar:
tempImg = self.removeBlackBorder(tempImg)
if RGBGama:
tempImg = cv2.cvtColor(tempImg, cv2.COLOR_BGR2RGB)
if preResize is not None:
tempImg = cv2.resize(tempImg, preResize, interpolation=cv2.INTER_AREA)#Cambio de resolucion
lista.append (np.mean(tempImg, axis=tuple(range(tempImg.ndim-1))))
if not (sPath/"train").exists():
imgPath = dumpImgs/f"Img{i}.npy"
np.save(imgPath, tempImg)
#parImgClass = (tempImg.astype(float), 0 if file.parent.name == "NOT_NVB" else 1)
parImgClass = (imgPath, 0 if file.parent.name == "NOT_NVB" else 1)
if file.parent.name == "NOT_NVB":
self.nNVB.append(parImgClass)
else:
self.yNVB.append(parImgClass)
if SegImage is not None:
for j, newImg in enumerate(self._sliceImage(tempImg, SegImage)):
imgPath = dumpImgs/f"Img{i}-{j}.npy"
np.save(imgPath, newImg)
#parImgClass = (tempImg.astype(float), 0 if file.parent.name == "NOT_NVB" else 1)
parImgClass = (imgPath, 0 if file.parent.name == "NOT_NVB" else 1)
if file.parent.name == "NOT_NVB":
self.nNVB.append(parImgClass)
else:
self.yNVB.append(parImgClass)
temp = np.asarray(lista)
self.mean = list(np.mean(temp, axis=tuple(range(temp.ndim-1))))
self.std = list(np.std(temp, axis=tuple(range(temp.ndim-1))))
if (sPath / "train").exists():
self.test_DataSet = None
self.train_DataSet = None
self.val_DataSet = None
self.nNVB = None
self.yNVB = None
self.SaveFold = sPath
return
if FoldSeed is not None:
np.random.seed(FoldSeed)
shuffle_nNVB = list(range(len(self.nNVB)))
shuffle_NVB = list(range(len(self.yNVB)))
np.random.shuffle(shuffle_nNVB)
np.random.shuffle(shuffle_NVB)
pasoN, pasoY = (dist['train'][0], dist['train'][1])
self.train_DataSet = self.nNVB[:pasoN] + self.yNVB[:pasoY]
nextPasoN, nextPasoY = (pasoN+dist['val'][0], pasoY+dist['val'][1])
self.val_DataSet = self.nNVB[pasoN:nextPasoN] + self.yNVB[pasoY:nextPasoY]
self.test_DataSet = self.nNVB[nextPasoN:] + self.yNVB[nextPasoY:]
if createFile:
if len(SavePath) == 0:
raise Exception("If createFile is True, SavePath must have a value")
n=0
for img, lb in self.train_DataSet:
trainPath = sPath / "train" / ("NOT_NVB" if lb == 0 else "NVB") / f"img{n}"
trainPath.parent.mkdir(parents=True, exist_ok=True)
np.save(trainPath, defs.load_file(img))
n += 1
n=0 if not Fold else n
for img, lb in self.val_DataSet:
trainPath = sPath / ("val" if not Fold else "train") / ("NOT_NVB" if lb == 0 else "NVB") / f"img{n}"
trainPath.parent.mkdir(parents=True, exist_ok=True)
np.save(trainPath, defs.load_file(img))
n += 1
n=0 if not Fold else n
for img, lb in self.test_DataSet:
trainPath = sPath / ("test" if not Fold else "train") / ("NOT_NVB" if lb == 0 else "NVB") / f"img{n}"
trainPath.parent.mkdir(parents=True, exist_ok=True)
np.save(trainPath, defs.load_file(img))
n += 1
self.test_DataSet = None
self.train_DataSet = None
self.val_DataSet = None
self.nNVB = None
self.yNVB = None
self.SaveFold = sPath
for f in dumpImgs.iterdir():
f.unlink()
dumpImgs.rmdir()
class DataloaderToDataModule(LightningDataModule):
"""Converts a set of dataloaders into a lightning datamodule.
Args:
train_dataloader: Training dataloader
val_dataloaders: Validation dataloader(s)
test_dataloaders: Test dataloader(s)
"""
def __init__(
self,
train_dataloader: DataLoader,
val_dataloaders: Union[DataLoader, Sequence[DataLoader]],
) -> None:
super().__init__()
self._train_dataloader = train_dataloader
self._val_dataloaders = val_dataloaders
def train_dataloader(self) -> DataLoader:
"""Return training dataloader."""
return self._train_dataloader
def val_dataloader(self) -> Union[DataLoader, Sequence[DataLoader]]:
"""Return validation dataloader(s)."""
return self._val_dataloaders
class RARP_DataModule(LightningDataModule):
def __init__(self,
KFold:bool=False,
num_fold:int=5,
shuffle: bool = False,
stratified: bool = False,
train_dataloader: DataLoader = None,
val_dataloader: DataLoader = None,
test_dataloader: DataLoader = None,
transforms = []) -> None:
super().__init__()
self.fold_index = 0
self.splits = None
self.Kfold_split = KFold
self.num_folds = num_fold
self.Folds = []
self.shuffle = shuffle
self.stratified = stratified
self.label_extractor = lambda batch: batch[1]
self.trainDL = train_dataloader
self.valDL = val_dataloader
self.testDL = test_dataloader
self.dataloader_settings = None
self.transforms = transforms
tempFoldsOrder = list(range(num_fold))
for _ in range(num_fold):
self.Folds.append(tempFoldsOrder)
tempFoldsOrder = tempFoldsOrder[1:] + tempFoldsOrder[:1]
def get_Current_Folds(self) -> list:
return self.Folds[self.fold_index]
def get_labels(self, dataloader: DataLoader) -> Optional[List]:
"""Try to extract the training labels (for classification problems) from the underlying training dataset."""
# Try to extract labels from the dataset through labels attribute
if hasattr(dataloader.dataset, "labels"):
return dataloader.dataset.labels.tolist()
# Else iterate and try to extract
try:
return torch.cat([self.label_extractor(batch) for batch in dataloader], dim=0).tolist()
except Exception:
return None
def _split(self, a, n):
k, m = divmod(len(a), n)
return (a[i*k+min(i, m):(i+1)*k+min(i+1, m)] for i in range(n))
def split_folds(self):
if self.splits is None:
self.splits = []
self.FoldDataset = self.trainDL.dataset
NO_NVB = []
NVB = []
for pos, label in enumerate(self.FoldDataset.targets):
if label==0:
NO_NVB.append(pos)
else:
NVB.append(pos)
if self.shuffle:
np.random.shuffle(NO_NVB)
np.random.shuffle(NVB)
NO_NVB_Folds = list (self._split(NO_NVB, self.num_folds))
NVB_Fols = list(self._split(NVB, self.num_folds))
setst = [
math.trunc(0.60 * self.num_folds), math.trunc(0.80 * self.num_folds), self.num_folds
]
tempArray = [] ##### FIX ... que puede hacer los grupos mal sin los folds son difrentes de 5
for pos, index in enumerate(self.Folds[self.fold_index]):
if pos < setst[0]:
tempArray += NO_NVB_Folds[index] + NVB_Fols[index]
elif pos < setst[1]:
self.splits.append(tempArray)
self.splits.append(NO_NVB_Folds[index] + NVB_Fols[index])
else:
self.splits.append(NO_NVB_Folds[index] + NVB_Fols[index])
def setup_folds(self):
if self.splits is None:
labels = None
if self.stratified:
labels = self.get_labels(self.trainDL)
if labels is None:
raise ValueError(
"Tried to extract labels for stratified K folds but failed."
" Make sure that the dataset of your train dataloader either"
" has an attribute `labels` or that `label_extractor` attribute"
" is initialized correctly"
)
splitter = StratifiedKFold(self.num_folds, shuffle=self.shuffle)
else:
splitter = KFold(self.num_folds, shuffle=self.shuffle)
self.FoldDataset = self.trainDL.dataset
self.splits = [split for split in splitter.split(range(len(self.FoldDataset)), y=labels)]
def train_dataloader(self) -> DataLoader:
if not self.Kfold_split:
return self.trainDL
else:
self.split_folds()
train_fold = RARP_Dataset(self.FoldDataset, self.splits[0], self.transforms[0])
return DataLoader(train_fold, shuffle=True, **self.dataloader_setting)
def val_dataloader(self) -> DataLoader:
if not self.Kfold_split:
return self.valDL
else:
self.split_folds()
val_fold = RARP_Dataset(self.FoldDataset, self.splits[1], self.transforms[1])
return DataLoader(val_fold, shuffle=False, **self.dataloader_setting)
def test_dataloader(self) -> DataLoader:
if not self.Kfold_split:
return self.testDL
else:
self.split_folds()
test_fold = RARP_Dataset(self.FoldDataset, self.splits[2], self.transforms[2])
return DataLoader(test_fold, shuffle=False, **self.dataloader_setting)
@property
def dataloader_setting(self) -> dict:
"""Return the settings of the train dataloader."""
if self.dataloader_settings is None:
orig_dl = self.trainDL
self.dataloader_settings = {
"batch_size": orig_dl.batch_size,
"num_workers": orig_dl.num_workers,
"collate_fn": orig_dl.collate_fn,
"pin_memory": orig_dl.pin_memory,
"drop_last": orig_dl.drop_last,
"timeout": orig_dl.timeout,
"worker_init_fn": orig_dl.worker_init_fn,
"prefetch_factor": orig_dl.prefetch_factor,
"persistent_workers": orig_dl.persistent_workers,
}
return self.dataloader_settings
class RARP_Dataset(Dataset):
def __init__(self, RARP_dataset, indices, transform=None) -> None:
super().__init__()
self.transform = transform
self.samples = Subset(RARP_dataset, indices)
self.targets = [s[1] for s in self.samples]
self.classes = ["NO_NVB", "NVB"] ## **** Hacer esta parte correctamente */***
def __len__(self):
return len(self.samples)
def __getitem__(self, index):
img, label = self.samples[index]
if self.transform is not None:
img = self.transform(img)
return (img, label)
class RARP_ChannelSwap(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
def forward(self, img):
# It is assumed that the input image is in RGB format.
Channels = [0, 1, 2]
np.random.shuffle(Channels)
return img[Channels]
class RARP_Invert(torch.nn.Module):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def forward(self, img):
return 255 - img
class RARP_DINO_Augmentation():
def __init__(
self,
GloblaCropsScale=(0.4, 1),
LocalCropsScale=(0.05, 0.4),
NumLocalCrops:int=8,
Size:int=224,
device = None,
mean:float = None,
std:float = None,
Tranform_0 = None
) -> None:
self.NumLocal_Crops= NumLocalCrops
self.globalCrop1 = torch.nn.Sequential(
torchvision.transforms.RandomResizedCrop(
Size,
scale=GloblaCropsScale,
antialias=True,
interpolation=torchvision.transforms.InterpolationMode.BICUBIC
),
torchvision.transforms.RandomHorizontalFlip(0.5),
torchvision.transforms.RandomErasing(0.2, value="random"),
#torchvision.transforms.RandomApply([
# RARP_ChannelSwap()
#]),
torchvision.transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2)),
torchvision.transforms.Normalize(mean, std)
).to(device)
self.globalCrop2 = torch.nn.Sequential(
torchvision.transforms.RandomResizedCrop(
Size,
scale=GloblaCropsScale,
antialias=True,
interpolation=torchvision.transforms.InterpolationMode.BICUBIC
),
torchvision.transforms.RandomHorizontalFlip(0.5),
torchvision.transforms.RandomErasing(0.8, value="random"),
#torchvision.transforms.RandomApply([
# RARP_ChannelSwap()
#]),
torchvision.transforms.RandomApply([
torchvision.transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2))
], 0.1),
torchvision.transforms.RandomApply([
RARP_Invert()
], 0.2),
torchvision.transforms.Normalize(mean, std)
).to(device)
self.local = torch.nn.Sequential(
torchvision.transforms.RandomResizedCrop(
Size,
scale=LocalCropsScale,
antialias=True,
interpolation=torchvision.transforms.InterpolationMode.BICUBIC
),
torchvision.transforms.RandomHorizontalFlip(0.5),
torchvision.transforms.RandomErasing(0.5, value="random"),
#torchvision.transforms.RandomApply([
# RARP_ChannelSwap()
#]),
torchvision.transforms.RandomApply([
torchvision.transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2))
], 0.5),
torchvision.transforms.Normalize(mean, std)
).to(device)
InitResize = (256,256)
self.classification = torch.nn.Sequential(
torchvision.transforms.Resize(
InitResize,
antialias=True,
interpolation=torchvision.transforms.InterpolationMode.BICUBIC
),
torchvision.transforms.CenterCrop(224),
torchvision.transforms.RandomAffine(
degrees=(-5, 5), scale=(0.9, 1.1), fill=5
),
torchvision.transforms.RandomHorizontalFlip(1.0),
torchvision.transforms.Normalize(mean, std),
).to(device) if Tranform_0 == None else Tranform_0
def __call__(self, img):
all_Crops = []
all_Crops.append(self.classification(img))
all_Crops.append(self.globalCrop1(img))
all_Crops.append(self.globalCrop2(img))
all_Crops.extend([self.local(img) for _ in range(self.NumLocal_Crops)])
return all_Crops
class RARP_DataSetType(Enum):
train = 0
val = 1
test = 2
class RandomVideoDataset(torch.utils.data.IterableDataset):
def __init__(self, root, epoch_size=None, frame_transform=None, video_transform=None, clip_len=16):
super().__init__()
self.samples = get_samples(root)
# Allow for temporal jittering
if epoch_size is None:
epoch_size = len(self.samples)
self.epoch_size = epoch_size
self.clip_len = clip_len
self.frame_transform = frame_transform
self.video_transform = video_transform
def __iter__(self):
for i in range(self.epoch_size):
# Get random sample
path, target = np.random.choice(self.samples)
# Get video object
vid = torchvision.io.VideoReader(path, "video")
metadata = vid.get_metadata()
video_frames = [] # video frame buffer
# Seek and return frames
max_seek = metadata["video"]['duration'][0] - (self.clip_len / metadata["video"]['fps'][0])
start = np.random.uniform(0., max_seek)
for frame in itertools.islice(vid.seek(start), self.clip_len):
video_frames.append(self.frame_transform(frame['data']))
current_pts = frame['pts']
# Stack it into a tensor
video = torch.stack(video_frames, 0)
if self.video_transform:
video = self.video_transform(video)
output = {
'path': path,
'video': video,
'target': target,
'start': start,
'end': current_pts}
yield output
class VideoDataset(Dataset):
def __init__(self, time_depth, mean, std, transform=None) -> None:
super().__init__()
class RARPDataset(Dataset):
def __init__(self, path_RARP_dataset:Path, path_RARP_Folds:Path=None, split=0, DataSetType:RARP_DataSetType="train", transform=None) -> None:
super().__init__()
self.samples = pd.read_csv(path_RARP_dataset, usecols=["id", "label", "class", "path"])
self.classes = self.samples["class"].unique().tolist()
self.targets = self.samples["label"].to_list()
self.IDsplit = None
self.transform = transform
if path_RARP_Folds is not None:
if split is None:
raise Exception("Is required the split index to do Folds")
arrFolds = np.load(path_RARP_Folds, allow_pickle=True)
self.IDsplit = np.array_split(arrFolds, len(arrFolds)/3)[split][DataSetType.value]
self.samples = self.samples.loc[self.samples["id"].isin(self.IDsplit)]
self.targets = self.samples["label"]
def __len__(self):
return self.samples.shape[0]
def __getitem__(self, index):
ID = index
if self.IDsplit is not None:
ID = self.IDsplit[index]
rs = self.samples.loc[self.samples["id"] == ID]
img = defs.load_file_tensor(rs["path"].item())
label = rs["label"].item()
if self.transform is not None:
img = self.transform(img)
return img, label
