|
|
import os |
|
|
import torch |
|
|
from torch.utils.data import Dataset |
|
|
from torchvision import transforms |
|
|
from PIL import Image, ImageOps |
|
|
import numpy as np |
|
|
import json |
|
|
import cv2 |
|
|
import albumentations as A |
|
|
from albumentations.pytorch import ToTensorV2 |
|
|
|
|
|
|
|
|
import logging |
|
|
|
|
|
logger = logging.getLogger() |
|
|
|
|
|
class ImageDataset(Dataset): |
|
|
def __init__(self, |
|
|
data_root, |
|
|
meta_file="", |
|
|
resize=240, |
|
|
mode="train", |
|
|
aug=False, |
|
|
with_fg_mask=False, |
|
|
test_class='None' |
|
|
): |
|
|
|
|
|
self.data_root = data_root |
|
|
self.resize = resize |
|
|
self.mode = mode |
|
|
self.test_class = test_class |
|
|
self.with_fg_mask = with_fg_mask |
|
|
self.aug = aug |
|
|
|
|
|
if isinstance(meta_file, str): |
|
|
meta_info = json.load(open(meta_file, 'r')) |
|
|
else: |
|
|
meta_info = meta_file |
|
|
|
|
|
self.data_list = [] |
|
|
if self.mode == "train": |
|
|
meta_info = meta_info[mode] |
|
|
for cls_name, data_list in meta_info.items(): |
|
|
self.data_list.extend(data_list) |
|
|
|
|
|
|
|
|
|
|
|
self.class_names = list(meta_info.keys()) |
|
|
else: |
|
|
meta_info = meta_info[mode][test_class] |
|
|
self.data_list.extend(meta_info) |
|
|
self.class_names = [test_class] |
|
|
|
|
|
|
|
|
self.resize_img_transform = transforms.Resize((self.resize, self.resize), interpolation=Image.BICUBIC) |
|
|
self.resize_mask_transform = transforms.Resize((self.resize, self.resize), interpolation=Image.NEAREST) |
|
|
self.aug_transform = A.Compose([ |
|
|
A.HorizontalFlip(p=0.2), |
|
|
A.VerticalFlip(p=0.2), |
|
|
A.ShiftScaleRotate(shift_limit=0.2, scale_limit=0, rotate_limit=0, p=0.2), |
|
|
|
|
|
ToTensorV2() |
|
|
]) |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
data = self.data_list[idx] |
|
|
img_path, mask_path, cls_name, anomaly = data["img_path"], data["mask_path"], data["cls_name"], data["anomaly"] |
|
|
|
|
|
img_path = os.path.join(self.data_root, img_path) |
|
|
mask_path = os.path.join(self.data_root, mask_path) |
|
|
if self.with_fg_mask: |
|
|
fg_mask_path = img_path.replace(self.data_root, "sam2_fg_mask") |
|
|
fg_mask_path = fg_mask_path[:-3] + "png" |
|
|
|
|
|
image = Image.open(img_path).convert('RGB') |
|
|
image = ImageOps.exif_transpose(image) |
|
|
image = self.resize_img_transform(image) |
|
|
|
|
|
if anomaly == 0: |
|
|
mask = Image.fromarray(np.zeros((self.resize, self.resize)), mode='L') |
|
|
else: |
|
|
mask = np.array(Image.open(mask_path).convert('L')) > 0 |
|
|
mask = Image.fromarray(mask.astype(np.uint8) * 255, mode='L') |
|
|
mask = self.resize_mask_transform(mask) |
|
|
|
|
|
if self.with_fg_mask: |
|
|
fg_mask = np.array(Image.open(fg_mask_path).convert('L')) > 0 |
|
|
fg_mask = Image.fromarray(fg_mask.astype(np.uint8) * 255, mode='L') |
|
|
fg_mask = self.resize_mask_transform(fg_mask) |
|
|
|
|
|
else: |
|
|
fg_mask = torch.zeros(1, self.resize, self.resize) |
|
|
|
|
|
if self.mode == "train" and self.aug: |
|
|
image = np.array(image).astype(np.float32) |
|
|
mask = np.array(mask) |
|
|
augmented = self.aug_transform(image=image, mask=mask) |
|
|
image = augmented['image'] |
|
|
mask = augmented['mask'] |
|
|
if self.with_fg_mask: |
|
|
fg_mask = np.array(fg_mask) |
|
|
fg_mask = self.aug_transform(mask=fg_mask)['mask'] |
|
|
|
|
|
else: |
|
|
image = transforms.ToTensor()(image) |
|
|
mask = transforms.ToTensor()(mask) |
|
|
if self.with_fg_mask: |
|
|
fg_mask = transforms.ToTensor()(fg_mask) |
|
|
|
|
|
return {"image": image, "mask": mask, "fg_mask": fg_mask, "cls_name": cls_name, "anomaly": anomaly} |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.data_list) |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
|
|
ds = ImageDataset(is_train=True) |
|
|
|
