bf in InferenceDataset in datasets.py

This commit is contained in:
Till Raab 2023-10-27 11:25:32 +02:00
parent 873758d461
commit f3e88078c4

View File

@ -24,15 +24,11 @@ class InferenceDataset(Dataset):
def __init__(self, dir_path): def __init__(self, dir_path):
self.dir_path = dir_path self.dir_path = dir_path
self.all_images = sorted(list(Path(self.dir_path).rglob(f'*.png'))) self.all_images = sorted(list(Path(self.dir_path).rglob(f'*.png')))
embed()
quit()
def __len__(self): def __len__(self):
return len(self.all_images) return len(self.all_images)
def __getitem__(self, idx): def __getitem__(self, idx):
image_name = self.all_images[idx] image_path = self.all_images[idx]
image_path = os.path.join(self.dir_path, image_name)
img = Image.open(image_path) img = Image.open(image_path)
img_tensor = F.to_tensor(img.convert('RGB')) img_tensor = F.to_tensor(img.convert('RGB'))