bf in InferenceDataset in datasets.py

2023-10-27 11:25:32 +02:00 · 2023-10-27 11:25:32 +02:00 · f3e88078c4
commit f3e88078c4
parent 873758d461
1 changed files with 1 additions and 5 deletions
--- a/datasets.py
+++ b/datasets.py
@ -24,15 +24,11 @@ class InferenceDataset(Dataset):
    def __init__(self, dir_path):
        self.dir_path = dir_path
        self.all_images = sorted(list(Path(self.dir_path).rglob(f'*.png')))
-        embed()
-        quit()
    def __len__(self):
        return len(self.all_images)

    def __getitem__(self, idx):
-        image_name = self.all_images[idx]
-        image_path = os.path.join(self.dir_path, image_name)
-
+        image_path = self.all_images[idx]
        img = Image.open(image_path)
        img_tensor = F.to_tensor(img.convert('RGB'))