from typing import List, Tuple
import os
import glob
import numpy as np
import pandas as pd
from PIL import Image
from scipy.ndimage.filters import gaussian_filter, median_filter, rank_filter
from torch.utils.data import Dataset
from torchvision import transforms
from utils.constants import Split, Columns, CropsColumns, ProbsColumns
from utils.paths import CROPS_DATASET, CROPS_PATH, COORDS_PATH, IMG_PATH, PROBS_DATASET, PROBS_PATH, HAADF_DATASET, PT_DATASET
class ImageClassificationDataset(Dataset):
def __init__(self, image_paths, image_labels, include_filename=False):
self.image_paths = image_paths
self.image_labels = image_labels
self.include_filename = include_filename
self.transform = transforms.Compose([
transforms.ToTensor()
# transforms.Normalize(mean=[0.5], std=[0.5])
])
def get_n_labels(self):
return len(set(self.image_labels))
def __len__(self):
return len(self.image_paths)
@staticmethod
def load_image(img_filename):
img = Image.open(img_filename)
np_img = np.asarray(img).astype(np.float32)
np_bg = median_filter(np_img, size=(40, 40))
np_clean = np_img - np_bg
np_normed = (np_clean - np_clean.min()) / (np_clean.max() - np_clean.min())
return np_normed
def __getitem__(self, idx):
img_path = self.image_paths[idx]
image = self.load_image(img_path)
image = self.transform(image)
label = self.image_labels[idx]
if self.include_filename:
return image, label, os.path.basename(img_path)
else:
return image, label
@staticmethod
def get_filenames_labels(split: Split) -> Tuple[List[str], List[int]]:
raise NotImplementedError
@classmethod
def train_dataset(cls, **kwargs):
filenames, labels = cls.get_filenames_labels(Split.TRAIN)
return cls(filenames, labels, **kwargs)
@classmethod
def val_dataset(cls, **kwargs):
filenames, labels = cls.get_filenames_labels(Split.VAL)
return cls(filenames, labels, **kwargs)
@classmethod
def test_dataset(cls, **kwargs):
filenames, labels = cls.get_filenames_labels(Split.TEST)
return cls(filenames, labels, **kwargs)
class HaadfDataset(ImageClassificationDataset):
@staticmethod
def get_filenames_labels(split: Split) -> Tuple[List[str], List[int]]:
df = pd.read_csv(HAADF_DATASET)
split_df = df[df[Columns.SPLIT] == split]
filenames = (IMG_PATH + os.sep + split_df[Columns.FILENAME]).to_list()
labels = (split_df[Columns.LABEL]).to_list()
return filenames, labels
class ImageDataset:
FILENAME_COL = "Filename"
SPLIT_COL = "Split"
RULER_UNITS = "Ruler Units"
def __init__(self, dataset_csv: str):
self.df = pd.read_csv(dataset_csv)
def iterate_data(self, split: Split):
df = self.df[self.df[self.SPLIT_COL] == split]
for idx, row in df.iterrows():
image_filename = os.path.join(IMG_PATH, row[self.FILENAME_COL])
yield image_filename
def get_ruler_units_by_img_name(self, name):
print(name)
return self.df[self.df[self.FILENAME_COL] == name][self.RULER_UNITS].values[0]
class CoordinatesDataset:
FILENAME_COL = "Filename"
COORDS_COL = "Coords"
SPLIT_COL = "Split"
def __init__(self, coord_image_csv: str):
self.df = pd.read_csv(coord_image_csv)
def iterate_data(self, split: Split):
df = self.df[self.df[self.SPLIT_COL] == split]
for idx, row in df.iterrows():
image_filename = os.path.join(IMG_PATH, row[self.FILENAME_COL])
if isinstance(row[self.COORDS_COL], str):
coords_filename = os.path.join(COORDS_PATH, row[self.COORDS_COL])
else:
coords_filename = None
yield image_filename, coords_filename
@staticmethod
def load_coordinates(label_filename: str) -> List[Tuple[int, int]]:
atom_coordinates = pd.read_csv(label_filename)
return list(zip(atom_coordinates['X'], atom_coordinates['Y']))
def split_length(self, split: Split):
df = self.df[self.df[self.SPLIT_COL] == split]
return len(df)
class HaadfCoordinates(CoordinatesDataset):
def __init__(self):
super().__init__(coord_image_csv=PT_DATASET)
class CropsDataset(ImageClassificationDataset):
@staticmethod
def get_filenames_labels(split: Split):
df = pd.read_csv(CROPS_DATASET)
split_df = df[df[CropsColumns.SPLIT] == split]
filenames = (CROPS_PATH + os.sep + split_df[CropsColumns.FILENAME]).to_list()
labels = (split_df[CropsColumns.LABEL]).to_list()
return filenames, labels
class CropsCustomDataset(ImageClassificationDataset):
@staticmethod
def get_filenames_labels(split: Split, crops_dataset: str, crops_path: str):
df = pd.read_csv(crops_dataset)
split_df = df[df[CropsColumns.SPLIT] == split]
filenames = (crops_path + os.sep + split_df[CropsColumns.FILENAME]).to_list()
labels = (split_df[CropsColumns.LABEL]).to_list()
return filenames, labels
class ProbsDataset(ImageClassificationDataset):
@staticmethod
def get_filenames_labels(split: Split):
df = pd.read_csv(PROBS_DATASET)
split_df = df[df[ProbsColumns.SPLIT] == split]
filenames = (PROBS_PATH + os.sep + split_df[ProbsColumns.FILENAME]).to_list()
labels = (split_df[ProbsColumns.LABEL]).to_list()
return filenames, labels
class SlidingCropDataset(Dataset):
def __init__(self, tif_filename, include_coords=True):
self.filename = tif_filename
self.include_coords = include_coords
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
])
self.n_labels = 2
self.step_size = 2
self.window_size = (21, 21)
self.loaded_crops = []
self.loaded_coords = []
self.load_crops()
def sliding_window(self, image):
# slide a window across the image
for x in range(0, image.shape[0] - self.window_size[0], self.step_size):
for y in range(0, image.shape[1] - self.window_size[1], self.step_size):
# yield the current window
center_x = x + ((self.window_size[0] - 1) // 2)
center_y = y + ((self.window_size[1] - 1) // 2)
yield center_x, center_y, image[x:x + self.window_size[0], y:y + self.window_size[1]]
@staticmethod
def load_image(img_filename):
img = Image.open(img_filename)
np_img = np.asarray(img).astype(np.float32)
np_bg = median_filter(np_img, size=(40, 40))
np_clean = np_img - np_bg
np_normed = (np_clean - np_clean.min()) / (np_clean.max() - np_clean.min())
return np_normed
def load_crops(self):
img = self.load_image(self.filename)
for x_center, y_center, img_crop in self.sliding_window(img):
self.loaded_crops.append(img_crop)
self.loaded_coords.append((x_center, y_center))
def get_n_labels(self):
return self.n_labels
def __len__(self):
return len(self.loaded_crops)
def __getitem__(self, idx):
crop = self.loaded_crops[idx]
x, y = self.loaded_coords[idx]
crop = self.transform(crop)
return crop, x, y
def get_image_path_without_coords(split: str or None = None):
coords_prefix_set = set()
for coords_name in os.listdir(COORDS_PATH):
coord_prefix = coords_name.split('_')[0]
coords_prefix_set.add(coord_prefix)
all_prefixes_set = set()
for tif_name in os.listdir(IMG_PATH):
coord_prefix = tif_name.split('_')[0]
all_prefixes_set.add(coord_prefix)
if split == Split.TRAIN:
missing_prefixes = coords_prefix_set
elif split == Split.TEST:
missing_prefixes = all_prefixes_set - coords_prefix_set
elif split is None:
missing_prefixes = all_prefixes_set
else:
raise ValueError
tif_filenames_list = []
labels_list = []
for prefix in missing_prefixes:
filename_matches = glob.glob(os.path.join(IMG_PATH, f'{prefix}_HAADF*NC*'))
if len(filename_matches) == 0:
continue
pos_filenames = [filename for filename in filename_matches if '_PtNC' in filename]
neg_filenames = [filename for filename in filename_matches if '_NC' in filename]
if len(pos_filenames) > 0:
pos_filename = sorted(pos_filenames)[-1]
tif_filenames_list.append(pos_filename)
labels_list.append(1)
if len(neg_filenames) > 0:
neg_filename = sorted(neg_filenames)[-1]
tif_filenames_list.append(neg_filename)
labels_list.append(0)
return tif_filenames_list, labels_list
if __name__ == "__main__":
filenames_list = get_image_path_without_coords()
filename = filenames_list[0]
dataset = SlidingCropDataset(filename)