import os from random import randint import numpy as np import torch from PIL import Image from torch.utils.data import Dataset from torchvision import transforms as T from torchvision.transforms import functional as F from typing import Callable import os import cv2 import pandas as pd from numbers import Number from typing import Container from collections import defaultdict from batchgenerators.utilities.file_and_folder_operations import * from collections import OrderedDict from torchvision.transforms import InterpolationMode import random from torch.utils.data import DataLoader from scipy.ndimage import label as clabel from tqdm import tqdm from torchvision.transforms import functional as Func import time def to_long_tensor(pic): # handle numpy array img = torch.from_numpy(np.array(pic, np.uint8)) # backward compatibility return img.long() def correct_dims(*images): corr_images = [] # print(images) for img in images: if len(img.shape) == 2: corr_images.append(np.expand_dims(img, axis=2)) else: corr_images.append(img) if len(corr_images) == 1: return corr_images[0] else: return corr_images def random_click(mask, class_id=1): indices = np.argwhere(mask != 0) indices[:, [0,1]] = indices[:, [1,0]] point_label = 1 if len(indices) == 0: point_label = 0 indices = np.argwhere(mask == 0) indices[:, [0,1]] = indices[:, [1,0]] pt = indices[np.random.randint(len(indices))] return pt[np.newaxis, :], [point_label] def fixed_click(mask, class_id=1): indices = np.argwhere(mask != 0) indices[:, [0,1]] = indices[:, [1,0]] point_label = 1 if len(indices) == 0: point_label = 0 indices = np.argwhere(mask == 0) indices[:, [0,1]] = indices[:, [1,0]] pt = indices[len(indices)//2] return pt[np.newaxis, :], [point_label] def random_clicks(mask, class_id = 1, prompts_number=10): indices = np.argwhere(mask == class_id) indices[:, [0,1]] = indices[:, [1,0]] point_label = 1 if len(indices) == 0: point_label = 0 indices = np.argwhere(mask != class_id) indices[:, [0,1]] = indices[:, [1,0]] pt_index = np.random.randint(len(indices), size=prompts_number) pt = indices[pt_index] point_label = np.repeat(point_label, prompts_number) return pt, point_label def pos_neg_clicks(mask, class_id=1, pos_prompt_number=5, neg_prompt_number=5): pos_indices = np.argwhere(mask == class_id) pos_indices[:, [0,1]] = pos_indices[:, [1,0]] pos_prompt_indices = np.random.randint(len(pos_indices), size=pos_prompt_number) pos_prompt = pos_indices[pos_prompt_indices] pos_label = np.repeat(1, pos_prompt_number) neg_indices = np.argwhere(mask != class_id) neg_indices[:, [0,1]] = neg_indices[:, [1,0]] neg_prompt_indices = np.random.randint(len(neg_indices), size=neg_prompt_number) neg_prompt = neg_indices[neg_prompt_indices] neg_label = np.repeat(0, neg_prompt_number) pt = np.vstack((pos_prompt, neg_prompt)) point_label = np.hstack((pos_label, neg_label)) return pt, point_label def random_bbox(mask, class_id=1, img_size=256): # return box = np.array([x1, y1, x2, y2]) indices = np.argwhere(mask == class_id) # Y X indices[:, [0,1]] = indices[:, [1,0]] # x, y if indices.shape[0] ==0: return np.array([-1, -1, img_size, img_size]) shiftw = randint(-int(0.9*img_size), int(1.1*img_size)) shifth = randint(-int(0.9*img_size), int(1.1*img_size)) shiftx = randint(-int(0.05*img_size), int(0.05*img_size)) shifty = randint(-int(0.05*img_size), int(0.05*img_size)) minx = np.min(indices[:, 0]) maxx = np.max(indices[:, 0]) miny = np.min(indices[:, 1]) maxy = np.max(indices[:, 1]) new_centerx = (minx + maxx)//2 + shiftx new_centery = (miny + maxy)//2 + shifty minx = np.max([new_centerx-shiftw//2, 0]) maxx = np.min([new_centerx+shiftw//2, img_size-1]) miny = np.max([new_centery-shifth//2, 0]) maxy = np.min([new_centery+shifth//2, img_size-1]) return np.array([minx, miny, maxx, maxy]) def fixed_bbox(mask, class_id = 1, img_size=256): indices = np.argwhere(mask != 0) # indices = np.argwhere(mask == class_id) # Y X (0, 1) indices[:, [0,1]] = indices[:, [1,0]] if indices.shape[0] ==0: return np.array([-1, -1, img_size, img_size]) minx = np.min(indices[:, 0]) maxx = np.max(indices[:, 0]) miny = np.min(indices[:, 1]) maxy = np.max(indices[:, 1]) return np.array([minx, miny, maxx, maxy]) class Transform2D_BCIHM: """ Performs augmentation on image and mask when called. Due to the randomness of augmentation transforms, it is not enough to simply apply the same Transform from torchvision on the image and mask separetely. Doing this will result in messing up the ground truth mask. To circumvent this problem, this class can be used, which will take care of the problems above. Args: crop: tuple describing the size of the random crop. If bool(crop) evaluates to False, no crop will be taken. p_flip: float, the probability of performing a random horizontal flip. color_jitter_params: tuple describing the parameters of torchvision.transforms.ColorJitter. If bool(color_jitter_params) evaluates to false, no color jitter transformation will be used. p_random_affine: float, the probability of performing a random affine transform using torchvision.transforms.RandomAffine. long_mask: bool, if True, returns the mask as LongTensor in label-encoded format. """ def __init__(self, mode='train', img_size=256, low_img_size=256, ori_size=256, crop=(32, 32), p_flip=0.5, p_rota=0.5, p_scale=0.0, p_gaussn=1.0, p_contr=0.0, p_gama=0.0, p_distor=0.0, color_jitter_params=(0.1, 0.1, 0.1, 0.1), p_random_affine=0, long_mask=False): self.mode = mode self.crop = crop self.p_flip = p_flip self.p_rota = p_rota self.p_scale = p_scale self.p_gaussn = p_gaussn self.p_gama = p_gama self.p_contr = p_contr self.p_distortion = p_distor self.img_size = img_size self.color_jitter_params = color_jitter_params if color_jitter_params: self.color_tf = T.ColorJitter(*color_jitter_params) self.p_random_affine = p_random_affine self.long_mask = long_mask self.low_img_size = low_img_size self.ori_size = ori_size def __call__(self, image, mask): # transforming to tensor image, mask = F.to_tensor(image), F.to_tensor(mask) # if self.mode == 'train': # # random horizontal flip # if np.random.rand() < self.p_flip: # image, mask = F.hflip(image), F.hflip(mask) # # random rotation # if np.random.rand() < self.p_rota: # angle = T.RandomRotation.get_params((-30, 30)) # image, mask = F.rotate(image, angle), F.rotate(mask, angle) # # random add gaussian noise # if np.random.rand() < self.p_gaussn: # image, mask = image.cpu().numpy().transpose(1,2,0), image.cpu().numpy().transpose(1,2,0) # ns = np.random.randint(3, 15) # noise = np.random.normal(loc=0, scale=1, size=(512, 512, 1)) * ns # noise = noise.astype(int) # image = np.array(image) + noise # image, mask = F.to_tensor(image), F.to_tensor(mask) # else: # pass # transforming to tensor image, mask = F.resize(image, (self.img_size, self.img_size), InterpolationMode.BILINEAR), F.resize(mask, (self.ori_size, self.ori_size), InterpolationMode.NEAREST) low_mask = F.resize(mask, (self.low_img_size, self.low_img_size), InterpolationMode.NEAREST) image = (image - image.min()) / (image.max() - image.min()) return image, mask, low_mask class BCIHM(Dataset): def __init__(self, dataset_path: str, split='train', joint_transform: Callable = None, fold=0, img_size=256, prompt = "click", class_id=1, one_hot_mask: int = False) -> None: self.fold = fold self.dataset_path = dataset_path self.one_hot_mask = one_hot_mask self.split = split # id_list_file = os.path.join('./dataset/excel', 'BCIHM.csv') id_list_file = '/data/wxh/Medical/tmz/metrics/brain_bleed/SAMIHS/BHSD/bhsd_2d_index.csv' df = pd.read_csv(id_list_file, encoding='gbk') # id_list_file = os.path.join(dataset_path, 'MainPatient/{0}.txt'.format(split)) if self.split == 'train': self.img_list = [name for id, name in enumerate(df['img']) if df['fold'][id] != self.fold and df['label'][id] > 0] self.gt_list = [label for id, label in enumerate(df['gt']) if df['fold'][id] != self.fold and df['label'][id] > 0] elif self.split == 'val': self.img_list = [name for id, name in enumerate(df['img']) if df['fold'][id] == self.fold] self.gt_list = [name for id, name in enumerate(df['gt']) if df['fold'][id] == self.fold] elif self.split == 'test': self.img_list = [name for id, name in enumerate(df['img']) if df['fold'][id] == self.fold] self.gt_list = [name for id, name in enumerate(df['gt']) if df['fold'][id] == self.fold] # self.ids = [id_.strip() for id_ in open(id_list_file)] self.prompt = prompt self.img_size = img_size self.class_id = class_id self.classes = 2 if joint_transform: self.joint_transform = joint_transform else: to_tensor = T.ToTensor() self.joint_transform = lambda x, y: (to_tensor(x), to_tensor(y)) def __len__(self): return len(self.img_list) def __getitem__(self, i): """Get the images""" name = self.img_list[i] img_path = os.path.join(self.dataset_path, name) mask_name = self.gt_list[i] msk_path = os.path.join(self.dataset_path, mask_name) image = np.load(img_path) mask = np.load(msk_path) class_id = 1 # fixed since only one class of foreground mask[mask > 0] = 1 image = np.clip(image, np.percentile(image, 0.05), np.percentile(image, 99.5)).astype(np.int16) mask = mask.astype(np.uint8) image, mask = correct_dims(image, mask) if self.joint_transform: image, mask, low_mask = self.joint_transform(image, mask) mask, low_mask = mask.squeeze(0), low_mask.squeeze(0) if self.one_hot_mask: assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2])).scatter_(0, mask.long(), 1) # --------- make the point prompt ---------- if self.prompt == 'click': point_label = 1 if 'train' in self.split: pt, point_label = random_click(np.array(mask), class_id) bbox = random_bbox(np.array(mask), class_id, self.img_size) else: pt, point_label = fixed_click(np.array(mask), class_id) bbox = fixed_bbox(np.array(mask), class_id, self.img_size) pt = pt * self.img_size / 512 mask[mask!=0] = 1 mask[mask!=1] = 0 low_mask[low_mask!=0] = 1 low_mask[low_mask!=1] = 0 point_labels = np.array(point_label) if self.one_hot_mask: assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2])).scatter_(0, mask.long(), 1) low_mask = low_mask.unsqueeze(0) mask = mask.unsqueeze(0) bbox = bbox * self.img_size / 512 return { 'image': image, 'label': mask, 'p_label': point_labels, 'pt': pt, 'bbox': bbox, 'low_mask':low_mask, 'image_name': name.split('/')[-1].split('.')[0] + '.png', 'class_id': class_id, } class Transform2D_Instance: """ Performs augmentation on image and mask when called. Due to the randomness of augmentation transforms, it is not enough to simply apply the same Transform from torchvision on the image and mask separetely. Doing this will result in messing up the ground truth mask. To circumvent this problem, this class can be used, which will take care of the problems above. Args: crop: tuple describing the size of the random crop. If bool(crop) evaluates to False, no crop will be taken. p_flip: float, the probability of performing a random horizontal flip. color_jitter_params: tuple describing the parameters of torchvision.transforms.ColorJitter. If bool(color_jitter_params) evaluates to false, no color jitter transformation will be used. p_random_affine: float, the probability of performing a random affine transform using torchvision.transforms.RandomAffine. long_mask: bool, if True, returns the mask as LongTensor in label-encoded format. """ def __init__(self, img_size=256, low_img_size=256, ori_size=256, crop=(32, 32), p_flip=0.0, p_rota=0.0, p_scale=0.0, p_gaussn=0.0, p_contr=0.0, p_gama=0.0, p_distor=0.0, color_jitter_params=(0.1, 0.1, 0.1, 0.1), p_random_affine=0, long_mask=False): self.crop = crop self.p_flip = p_flip self.p_rota = p_rota self.p_scale = p_scale self.p_gaussn = p_gaussn self.p_gama = p_gama self.p_contr = p_contr self.p_distortion = p_distor self.img_size = img_size self.color_jitter_params = color_jitter_params if color_jitter_params: self.color_tf = T.ColorJitter(*color_jitter_params) self.p_random_affine = p_random_affine self.long_mask = long_mask self.low_img_size = low_img_size self.ori_size = ori_size def __call__(self, image, mask): # transforming to tensor image, mask = F.to_tensor(image), F.to_tensor(mask) # if self.mode == 'train': # # random horizontal flip # if np.random.rand() < self.p_flip: # image, mask = F.hflip(image), F.hflip(mask) # # random rotation # if np.random.rand() < self.p_rota: # angle = T.RandomRotation.get_params((-30, 30)) # image, mask = F.rotate(image, angle), F.rotate(mask, angle) # # random add gaussian noise # if np.random.rand() < self.p_gaussn: # image, mask = image.cpu().numpy().transpose(1,2,0), image.cpu().numpy().transpose(1,2,0) # ns = np.random.randint(3, 15) # noise = np.random.normal(loc=0, scale=1, size=(512, 512, 1)) * ns # noise = noise.astype(int) # image = np.array(image) + noise # image, mask = F.to_tensor(image), F.to_tensor(mask) # else: # pass # transforming to tensor image, mask = F.resize(image, (self.img_size, self.img_size), InterpolationMode.BILINEAR), F.resize(mask, (self.ori_size, self.ori_size), InterpolationMode.NEAREST) low_mask = F.resize(mask, (self.low_img_size, self.low_img_size), InterpolationMode.NEAREST) image = (image - image.min()) / (image.max() - image.min()) return image, mask, low_mask class Instance(Dataset): def __init__(self, dataset_path: str, split='train', joint_transform: Callable = None, fold=0, img_size=256, prompt = "click", class_id=1, one_hot_mask: int = False) -> None: self.fold = fold self.dataset_path = dataset_path self.one_hot_mask = one_hot_mask self.split = split id_list_file = os.path.join('./dataset/excel', 'Instance.csv') df = pd.read_csv(id_list_file, encoding='gbk') # id_list_file = os.path.join(dataset_path, 'MainPatient/{0}.txt'.format(split)) if self.split == 'train': self.img_list = [name for id, name in enumerate(df['img']) if df['fold'][id] != self.fold and df['label'][id] > 0] self.gt_list = [label for id, label in enumerate(df['gt']) if df['fold'][id] != self.fold and df['label'][id] > 0] elif self.split == 'val': self.img_list = [name for id, name in enumerate(df['img']) if df['fold'][id] == self.fold] self.gt_list = [name for id, name in enumerate(df['gt']) if df['fold'][id] == self.fold] elif self.split == 'test': self.img_list = [name for id, name in enumerate(df['img']) if df['fold'][id] == self.fold] self.gt_list = [name for id, name in enumerate(df['gt']) if df['fold'][id] == self.fold] # self.ids = [id_.strip() for id_ in open(id_list_file)] self.prompt = prompt self.img_size = img_size self.class_id = class_id self.classes = 2 if joint_transform: self.joint_transform = joint_transform else: to_tensor = T.ToTensor() self.joint_transform = lambda x, y: (to_tensor(x), to_tensor(y)) def __len__(self): return len(self.img_list) def __getitem__(self, i): """Get the images""" name = self.img_list[i] img_path = os.path.join(self.dataset_path, name) mask_name = self.gt_list[i] msk_path = os.path.join(self.dataset_path, mask_name) image = np.load(img_path) mask = np.load(msk_path) class_id = 1 # fixed since only one class of foreground mask[mask > 0] = 1 image = np.clip(image, np.percentile(image, 0.05), np.percentile(image, 99.5)).astype(np.int16) mask = mask.astype(np.uint8) image, mask = correct_dims(image, mask) if self.joint_transform: image, mask, low_mask = self.joint_transform(image, mask) mask, low_mask = mask.squeeze(0), low_mask.squeeze(0) if self.one_hot_mask: assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2])).scatter_(0, mask.long(), 1) # --------- make the point prompt ----------------- if self.prompt == 'click': point_label = 1 if 'train' in self.split: pt, point_label = random_click(np.array(mask), class_id) bbox = random_bbox(np.array(mask), class_id, self.img_size) else: pt, point_label = fixed_click(np.array(mask), class_id) bbox = fixed_bbox(np.array(mask), class_id, self.img_size) pt = pt * self.img_size / 512 mask[mask!=0] = 1 mask[mask!=1] = 0 low_mask[low_mask!=0] = 1 low_mask[low_mask!=1] = 0 point_labels = np.array(point_label) if self.one_hot_mask: assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2])).scatter_(0, mask.long(), 1) low_mask = low_mask.unsqueeze(0) mask = mask.unsqueeze(0) bbox = bbox * self.img_size / 512 return { 'image': image, 'label': mask, 'p_label': point_labels, 'pt': pt, 'bbox': bbox, 'low_mask':low_mask, 'image_name': name.split('/')[-1].split('.')[0] + '.png', 'class_id': class_id, } class Transform2D_Extended: def __init__(self, mode='train', img_size=256, low_img_size=256, ori_size=256, crop=(32, 32), p_flip=0.5, p_rota=0.5, p_scale=0.0, p_gaussn=1.0, p_contr=0.0, p_gama=0.0, p_distor=0.0, color_jitter_params=(0.1, 0.1, 0.1, 0.1), p_random_affine=0, long_mask=False): self.mode = mode self.crop = crop self.p_flip = p_flip self.p_rota = p_rota self.p_scale = p_scale self.p_gaussn = p_gaussn self.p_gama = p_gama self.p_contr = p_contr self.p_distortion = p_distor self.img_size = img_size self.color_jitter_params = color_jitter_params if color_jitter_params: self.color_tf = T.ColorJitter(*color_jitter_params) self.p_random_affine = p_random_affine self.long_mask = long_mask self.low_img_size = low_img_size self.ori_size = ori_size self._warned_flat = False # 仅打印一次 def __call__(self, image, mask): # to tensor image, mask = F.to_tensor(image), F.to_tensor(mask) # resize image = F.resize(image, (self.img_size, self.img_size), InterpolationMode.BILINEAR) mask = F.resize(mask, (self.ori_size, self.ori_size), InterpolationMode.NEAREST) low_mask = F.resize(mask, (self.low_img_size, self.low_img_size), InterpolationMode.NEAREST) # ---- SAFE min-max norm（避免分母为0 -> NaN）---- imn = image.min() imx = image.max() den = (imx - imn).clamp(min=1e-6) # 关键：加下限 if not self._warned_flat and (imx - imn) < 1e-6: # 只提示一次，说明存在“近乎常数图”的样本（多半是纯背景切片） # print("[Transform2D_Extended] Warning: near-constant slice encountered; " # "min==max normalization guarded to avoid NaN.") self._warned_flat = True image = (image - imn) / den # --------------------------------------------- # 最后一层保险（理论上不会触发） image = torch.nan_to_num(image, nan=0.0, posinf=0.0, neginf=0.0) return image, mask, low_mask class Extended(Dataset): def __init__(self, dataset_path_list: list, split='train', joint_transform: Callable = None, fold=0, img_size=256, prompt="click", class_id=1, one_hot_mask: int = False, pos_keep: float = 1.0, neg_keep: float = 0.01, seed: int = 1234) -> None: self.fold = fold self.dataset_path_list = dataset_path_list self.one_hot_mask = one_hot_mask self.split = split self.pos_keep = float(pos_keep) self.neg_keep = float(neg_keep) self.rng = np.random.default_rng(seed) id_list_file = '2d_index.csv' self.img_list, self.gt_list = [], [] for dataset_path in dataset_path_list: df = pd.read_csv(os.path.join(dataset_path, id_list_file), encoding='gbk') fold_col = df['fold'].to_numpy() is_pos = (df['label'].to_numpy() > 0) if self.split == 'train': in_fold = (fold_col != self.fold) # 随机保留比例（正样本用 pos_keep，负样本用 neg_keep） keep_pos = self.rng.random(len(df)) < self.pos_keep keep_neg = self.rng.random(len(df)) < self.neg_keep if 'Negative' in dataset_path: print(f'Negative: {len(df)}') keep_neg = self.rng.random(len(df)) < 0.6 else: keep_neg = self.rng.random(len(df)) < self.neg_keep keep = in_fold & ((is_pos & keep_pos) | (~is_pos & keep_neg)) elif self.split in ('val', 'test'): keep = (fold_col == self.fold) else: raise ValueError(f'Unknown split: {self.split}') idx = np.nonzero(keep)[0] img_add = [os.path.join(dataset_path, df['img'].iat[i]) for i in idx] gt_add = [os.path.join(dataset_path, df['gt' ].iat[i]) for i in idx] self.img_list += img_add self.gt_list += gt_add # 打印非累计统计（本数据源本次新增） pos_cnt = int((df.loc[idx, 'label'] > 0).sum()) neg_cnt = int((df.loc[idx, 'label'] == 0).sum()) print(f"{self.split} in {dataset_path}: +{len(idx)} (pos {pos_cnt}, neg {neg_cnt}); total {len(self.img_list)}") self.prompt = prompt self.img_size = img_size self.class_id = class_id self.classes = 2 if joint_transform: self.joint_transform = joint_transform else: to_tensor = T.ToTensor() self.joint_transform = lambda x, y: (to_tensor(x), to_tensor(y)) print(f"Dataset {self.split} size: {len(self.img_list)}") def __len__(self): return len(self.img_list) def __getitem__(self, i): """Get the images""" img_path = self.img_list[i] # img_path = os.path.join(self.dataset_path, name) msk_path = self.gt_list[i] # msk_path = os.path.join(self.dataset_path, mask_name) image = np.load(img_path) mask = np.load(msk_path) class_id = 1 # fixed since only one class of foreground mask[mask > 0] = 1 image = np.clip(image, np.percentile(image, 0.05), np.percentile(image, 99.5)).astype(np.int16) mask = mask.astype(np.uint8) image, mask = correct_dims(image, mask) if self.joint_transform: image, mask, low_mask = self.joint_transform(image, mask) mask, low_mask = mask.squeeze(0), low_mask.squeeze(0) if self.one_hot_mask: assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2])).scatter_(0, mask.long(), 1) # --------- make the point prompt ---------- if self.prompt == 'click': point_label = 1 if 'train' in self.split: pt, point_label = random_click(np.array(mask), class_id) bbox = random_bbox(np.array(mask), class_id, self.img_size) else: pt, point_label = fixed_click(np.array(mask), class_id) bbox = fixed_bbox(np.array(mask), class_id, self.img_size) pt = pt * self.img_size / 512 mask[mask!=0] = 1 mask[mask!=1] = 0 low_mask[low_mask!=0] = 1 low_mask[low_mask!=1] = 0 point_labels = np.array(point_label) if self.one_hot_mask: assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2])).scatter_(0, mask.long(), 1) low_mask = low_mask.unsqueeze(0) mask = mask.unsqueeze(0) bbox = bbox * self.img_size / 512 return { 'image': image, 'label': mask, 'p_label': point_labels, 'pt': pt, 'bbox': bbox, 'low_mask':low_mask, 'image_name': img_path.split('/')[-1].split('.')[0] + '.png', 'class_id': class_id, } # class Extended(Dataset): # def __init__(self, dataset_path_list: list, split='train', joint_transform: Callable = None, fold=0, # img_size=256, prompt="click", class_id=1, one_hot_mask: int = False, # pos_keep: float = 1.0, neg_keep: float = 0.2, seed: int = 1234, # # === RUNTIME NAN CHECK === # runtime_nan_scan: bool = True, # runtime_nan_log_dir: str = "./nan_reports_runtime", # sanitize_nan: bool = True) -> None: # self.fold = fold # self.dataset_path_list = dataset_path_list # self.one_hot_mask = one_hot_mask # self.split = split # self.pos_keep = float(pos_keep) # self.neg_keep = float(neg_keep) # self.rng = np.random.default_rng(seed) # # === RUNTIME NAN CHECK === # self.runtime_nan_scan = bool(runtime_nan_scan) # self.sanitize_nan = bool(sanitize_nan) # os.makedirs(runtime_nan_log_dir, exist_ok=True) # ts = time.strftime("%Y%m%d-%H%M%S") # self._runtime_log_path = os.path.join(runtime_nan_log_dir, f"nan_runtime_{split}_{ts}.csv") # self._runtime_log_header_written = False # id_list_file = '2d_index.csv' # self.img_list, self.gt_list = [], [] # for dataset_path in dataset_path_list: # df = pd.read_csv(os.path.join(dataset_path, id_list_file), encoding='gbk') # fold_col = df['fold'].to_numpy() # is_pos = (df['label'].to_numpy() > 0) # if self.split == 'train': # in_fold = (fold_col != self.fold) # keep_pos = self.rng.random(len(df)) < self.pos_keep # keep_neg = self.rng.random(len(df)) < self.neg_keep # keep = in_fold & ((is_pos & keep_pos) | (~is_pos & keep_neg)) # elif self.split in ('val', 'test'): # keep = (fold_col == self.fold) # else: # raise ValueError(f'Unknown split: {self.split}') # idx = np.nonzero(keep)[0] # img_add = [os.path.join(dataset_path, df['img'].iat[i]) for i in idx] # gt_add = [os.path.join(dataset_path, df['gt' ].iat[i]) for i in idx] # self.img_list += img_add # self.gt_list += gt_add # pos_cnt = int((df.loc[idx, 'label'] > 0).sum()) # neg_cnt = int((df.loc[idx, 'label'] == 0).sum()) # print(f"{self.split} in {dataset_path}: +{len(idx)} (pos {pos_cnt}, neg {neg_cnt}); total {len(self.img_list)}") # self.prompt = prompt # self.img_size = img_size # self.class_id = class_id # self.classes = 2 # if joint_transform: # self.joint_transform = joint_transform # else: # to_tensor = T.ToTensor() # self.joint_transform = lambda x, y: (to_tensor(x), to_tensor(y)) # print(f"Dataset {self.split} size: {len(self.img_list)}") # # === RUNTIME NAN CHECK === helper: safe stats for np/torch # def _stats(self, arr): # try: # if isinstance(arr, torch.Tensor): # return float(torch.nanmin(arr).item()), float(torch.nanmax(arr).item()), float(torch.nanmean(arr).item()) # a = np.asarray(arr) # return float(np.nanmin(a)), float(np.nanmax(a)), float(np.nanmean(a)) # except Exception: # return float('nan'), float('nan'), float('nan') # # === RUNTIME NAN CHECK === helper: append CSV # def _log_nan(self, idx, img_path, msk_path, stage, which, arr): # if not self.runtime_nan_scan: # return # header = not os.path.exists(self._runtime_log_path) or not self._runtime_log_header_written # try: # with open(self._runtime_log_path, "a", newline="", encoding="utf-8") as f: # if header: # f.write("index,split,stage,which,img_path,gt_path,min,max,mean\n") # self._runtime_log_header_written = True # mn, mx, me = self._stats(arr) # f.write(f"{idx},{self.split},{stage},{which},{img_path},{msk_path},{mn},{mx},{me}\n") # except Exception as e: # print(f"[RUNTIME-NAN] failed to write log: {e}") # # === RUNTIME NAN CHECK === helper: nan_to_num # def _sanitize(self, arr): # if not self.sanitize_nan: # return arr # if isinstance(arr, torch.Tensor): # return torch.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0) # a = np.asarray(arr) # return np.nan_to_num(a, nan=0.0, posinf=0.0, neginf=0.0) # def __len__(self): # return len(self.img_list) # def __getitem__(self, i): # img_path = self.img_list[i] # msk_path = self.gt_list[i] # # ---- Stage A: load ---- # image = np.load(img_path) # mask = np.load(msk_path) # if self.runtime_nan_scan: # if not np.isfinite(image).all(): # self._log_nan(i, img_path, msk_path, "A_load", "image", image) # image = self._sanitize(image) # if not np.isfinite(mask).all(): # self._log_nan(i, img_path, msk_path, "A_load", "mask", mask) # mask = self._sanitize(mask) # class_id = 1 # mask[mask > 0] = 1 # # ---- Stage B: clip & dims ---- # image = np.clip(image, np.percentile(image, 0.05), np.percentile(image, 99.5)).astype(np.int16) # mask = mask.astype(np.uint8) # image, mask = correct_dims(image, mask) # if self.runtime_nan_scan: # if isinstance(image, np.ndarray) and not np.isfinite(image).all(): # self._log_nan(i, img_path, msk_path, "B_clip_dims", "image", image) # image = self._sanitize(image) # if isinstance(mask, np.ndarray) and not np.isfinite(mask).all(): # self._log_nan(i, img_path, msk_path, "B_clip_dims", "mask", mask) # mask = self._sanitize(mask) # # ---- Stage C: joint_transform ---- # if self.joint_transform: # image, mask, low_mask = self.joint_transform(image, mask) # # squeeze back to 2D # mask, low_mask = mask.squeeze(0), low_mask.squeeze(0) # else: # low_mask = torch.from_numpy(mask.copy()) # if self.runtime_nan_scan: # if isinstance(image, torch.Tensor) and not torch.isfinite(image).all(): # self._log_nan(i, img_path, msk_path, "C_transform", "image", image) # image = self._sanitize(image) # if isinstance(mask, torch.Tensor) and not torch.isfinite(mask).all(): # self._log_nan(i, img_path, msk_path, "C_transform", "mask", mask) # mask = self._sanitize(mask) # if isinstance(low_mask, torch.Tensor) and not torch.isfinite(low_mask).all(): # self._log_nan(i, img_path, msk_path, "C_transform", "low_mask", low_mask) # low_mask = self._sanitize(low_mask) # # ---- Stage D: one-hot (optional) ---- # if self.one_hot_mask: # assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' # mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2]), dtype=torch.float32) # mask = mask.scatter_(0, mask.long(), 1) # # ---- Stage E: prompts & binarize ---- # if self.prompt == 'click': # if 'train' in self.split: # pt, point_label = random_click(np.array(mask), class_id) # bbox = random_bbox(np.array(mask), class_id, self.img_size) # else: # pt, point_label = fixed_click(np.array(mask), class_id) # bbox = fixed_bbox(np.array(mask), class_id, self.img_size) # pt = pt * self.img_size / 512 # # ensure binary # mask = (mask != 0).to(mask.dtype) # low_mask = (low_mask != 0).to(low_mask.dtype) # point_labels = np.array(point_label) # if self.one_hot_mask: # assert self.one_hot_mask > 0, 'one_hot_mask must be nonnegative' # mask = torch.zeros((self.one_hot_mask, mask.shape[1], mask.shape[2]), dtype=torch.float32).scatter_(0, mask.long(), 1) # # ---- Stage F: final pack ---- # low_mask = low_mask.unsqueeze(0) # mask = mask.unsqueeze(0) # bbox = bbox * self.img_size / 512 # # === RUNTIME NAN CHECK === final guard # if self.runtime_nan_scan: # for stage_name, arr, which in [ # ("F_final", image, "image"), # ("F_final", mask, "mask"), # ("F_final", low_mask, "low_mask"), # ]: # if (isinstance(arr, torch.Tensor) and not torch.isfinite(arr).all()) or \ # (isinstance(arr, np.ndarray) and not np.isfinite(arr).all()): # self._log_nan(i, img_path, msk_path, stage_name, which, arr) # arr = self._sanitize(arr) # return { # 'image': image, # 'label': mask, # 'p_label': point_labels, # 'pt': pt, # 'bbox': bbox, # 'low_mask': low_mask, # 'image_name': img_path.split('/')[-1].split('.')[0] + '.png', # 'class_id': class_id, # } class Transform2D_Unlabeled: """ Performs augmentation on image and mask when called. Due to the randomness of augmentation transforms, it is not enough to simply apply the same Transform from torchvision on the image and mask separetely. Doing this will result in messing up the ground truth mask. To circumvent this problem, this class can be used, which will take care of the problems above. Args: crop: tuple describing the size of the random crop. If bool(crop) evaluates to False, no crop will be taken. p_flip: float, the probability of performing a random horizontal flip. color_jitter_params: tuple describing the parameters of torchvision.transforms.ColorJitter. If bool(color_jitter_params) evaluates to false, no color jitter transformation will be used. p_random_affine: float, the probability of performing a random affine transform using torchvision.transforms.RandomAffine. long_mask: bool, if True, returns the mask as LongTensor in label-encoded format. """ def __init__(self, img_size=256, low_img_size=256, ori_size=256, crop=(32, 32), p_flip=0.0, p_rota=0.0, p_scale=0.0, p_gaussn=0.0, p_contr=0.0, p_gama=0.0, p_distor=0.0, color_jitter_params=(0.1, 0.1, 0.1, 0.1), p_random_affine=0, long_mask=False): self.crop = crop self.p_flip = p_flip self.p_rota = p_rota self.p_scale = p_scale self.p_gaussn = p_gaussn self.p_gama = p_gama self.p_contr = p_contr self.p_distortion = p_distor self.img_size = img_size self.color_jitter_params = color_jitter_params if color_jitter_params: self.color_tf = T.ColorJitter(*color_jitter_params) self.p_random_affine = p_random_affine self.long_mask = long_mask self.low_img_size = low_img_size self.ori_size = ori_size def __call__(self, image): # transforming to tensor image = F.to_tensor(image) # transforming to tensor image= F.resize(image, (self.img_size, self.img_size), InterpolationMode.BILINEAR) image = (image - image.min()) / (image.max() - image.min()) return image class Unlabeled(Dataset): def __init__(self, dataset_path: str, fold=0, img_size=256, one_hot_mask: int = False) -> None: self.fold = fold self.dataset_path = dataset_path self.one_hot_mask = one_hot_mask id_list_file = os.path.join('./dataset/excel', 'Unlabeled.csv') df = pd.read_csv(id_list_file, encoding='gbk') # id_list_file = os.path.join(dataset_path, 'MainPatient/{0}.txt'.format(split)) self.img_list = [name for id, name in enumerate(df['img'])] def __len__(self): return len(self.img_list) def __getitem__(self, i): """Get the images""" name = self.img_list[i] img_path = os.path.join(self.dataset_path, name) # 加载 image（numpy） image = np.load(img_path) # shape: [H, W] # 强制 clip，防止极端值影响 image = np.clip(image, np.percentile(image, 0.05), np.percentile(image, 99.5)) # normalize to [0, 1] in float32 image = (image - image.min()) / (image.max() - image.min() + 1e-6) image = image.astype(np.float32) # resize to 1024 x 1024 using cv2 (BILINEAR) image = cv2.resize(image, (1024, 1024), interpolation=cv2.INTER_LINEAR) # 转为 tensor，shape [1, H, W] image = torch.from_numpy(image).unsqueeze(0) # [1, 1024, 1024] return { 'image': image, 'image_name': name.split('/')[-1].split('.')[0] + '.png', } class Logger: def __init__(self, verbose=False): self.logs = defaultdict(list) self.verbose = verbose def log(self, logs): for key, value in logs.items(): self.logs[key].append(value) if self.verbose: print(logs) def get_logs(self): return self.logs def to_csv(self, path): pd.DataFrame(self.logs).to_csv(path, index=None) if __name__ == '__main__': tf_val = Transform2D_BCIHM(mode='train', img_size=1024, low_img_size=256, ori_size=512, crop=None, p_flip=1, color_jitter_params=None, long_mask=True) val_dataset = BCIHM("/data/openData_Med/BCIHM/", "test", tf_val, img_size=1024, class_id=1) Idataset = Instance("/data/openData_Med/Instance/", "test", tf_val, img_size=1024, class_id=1) valloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=True) with tqdm(total=len(valloader), desc='Validation round', unit='batch', leave=False) as pbar: for batch_idx, (datapack) in enumerate(valloader): pass