WasteDataset.py

import torch
from torch.utils.data import Dataset
from PIL import Image
import os

class WasteDataset(Dataset):
    def __init__(self,dataframe,image_dir,transform=None):
        self.dataframe = dataframe
        self.image_dir = image_dir
        self.transform = transform

    def __len__(self):
        return len(self.dataframe)
    
    def __getitem__(self,idx):        
        # getting the image name
        imgName = self.dataframe.iloc[idx,0]
        
        #getting the label of the image
        imgLabel = self.dataframe.iloc[idx,1]

        # Path vers l'image
        imgPath = os.path.join(self.image_dir,imgName)

        # Opening de l'image
        openedImg = Image.open(imgPath).convert("RGB")

        if self.transform is not None:
            openedImg = self.transform(openedImg)

        label_tensor = torch.tensor(imgLabel,dtype=torch.long)
        return openedImg, label_tensor
    

# 1. Le Dataset de test reste le même
class WasteTestDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        self.image_dir = image_dir
        self.image_names = os.listdir(image_dir)
        self.transform = transform

    def __len__(self):
        return len(self.image_names)
    
    def __getitem__(self, idx):
        img_name = self.image_names[idx]
        img_path = os.path.join(self.image_dir, img_name)
        
        openedImg = Image.open(img_path).convert("RGB")
        if self.transform is not None:
            openedImg = self.transform(openedImg)
            
        return openedImg, img_name