fixed image generation viton hd

api/model/src/datasets/vitonhd.py

@@ -107,9 +107,15 @@ class VitonHDDataset(data.Dataset):
                 c_names.append(c_name)
                 dataroot_names.append(dataroot)
 
-        self.im_names = im_names
-        self.c_names = c_names
-        self.dataroot_names = dataroot_names
+        #self.im_names = im_names
+        #self.c_names = c_names
+        #self.dataroot_names = dataroot_names
+
+        # Filtra le immagini e i nomi dei vestiti in base a quelle presenti in captions_dict
+        self.im_names = [im for im in im_names if im.split('_')[0] in self.captions_dict]
+        self.c_names = [c_names[idx] for idx, im in enumerate(im_names) if im.split('_')[0] in self.captions_dict]
+        self.dataroot_names = [dataroot_names[idx] for idx, im in enumerate(im_names) if im.split('_')[0] in self.captions_dict]
+
 
     def __getitem__(self, index):
         """