Coverage for dqm/domain_gap/custom_datasets.py: 53%

1from torch.utils.data import Dataset 

2from PIL import Image 

3import os 

4 

5 

6# Définir une classe de Dataset personnalisée pour charger les images à partir du fichier .txt 

7class ImagesWithLabelsDataset(Dataset): 

8 def __init__(self, txt_file, transform=None): 

9 self.txt_file = txt_file 

10 self.transform = transform 

11 

12 # Lire le fichier .txt et stocker les chemins d'images et les labels dans des listes 

13 self.images = [] 

14 self.labels = [] 

15 with open(txt_file, "r") as file: 

16 lines = file.readlines() 

17 for line in lines: 

18 image_path, label = line.strip().split(" ") 

19 self.images.append(image_path) 

20 self.labels.append(int(label)) 

21 

22 def __len__(self): 

23 return len(self.images) 

24 

25 def __getitem__(self, idx): 

26 image_path = self.images[idx] 

27 label = self.labels[idx] 

28 image = Image.open(image_path).convert("RGB") # Ouvrir l'image en mode RGB 

29 if self.transform: 

30 image = self.transform(image) 

31 return image, label 

32 

33 

34class ImagesFromFolderDataset(Dataset): 

35 def __init__(self, folder_path, transform=None): 

36 self.folder_path = folder_path 

37 self.transform = transform 

38 self.images = os.listdir(self.folder_path) 

39 

40 def __len__(self): 

41 return len(self.images) 

42 

43 def __getitem__(self, idx): 

44 image = Image.open(os.path.join(self.folder_path, self.images[idx])) 

45 if self.transform: 

46 image = self.transform(image) 

47 return image