-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataset.py
52 lines (47 loc) · 1.68 KB
/
dataset.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import os
from glob import glob
import torch
import torch.utils.data
from PIL import Image
from torchvision import transforms
class MVTecDataset(torch.utils.data.Dataset):
def __init__(self, root, category, input_size, is_train=True):
self.image_transform = transforms.Compose(
[
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
if is_train:
self.image_files = glob(
os.path.join(root, category, "train", "good", "*.png")
)
else:
self.image_files = glob(os.path.join(root, category, "test", "*", "*.png"))
self.target_transform = transforms.Compose(
[
transforms.Resize(input_size),
transforms.ToTensor(),
]
)
self.is_train = is_train
def __getitem__(self, index):
image_file = self.image_files[index]
image = Image.open(image_file)
image = self.image_transform(image)
if self.is_train:
return image
else:
if os.path.dirname(image_file).endswith("good"):
target = torch.zeros([1, image.shape[-2], image.shape[-1]])
else:
target = Image.open(
image_file.replace("/test/", "/ground_truth/").replace(
".png", "_mask.png"
)
)
target = self.target_transform(target)
return image, target, image_file
def __len__(self):
return len(self.image_files)