data.py

from torch.utils.data import Dataset
import utils
from torchvision import transforms
import os
from PIL import Image
class ODData(Dataset):
    def __init__(self,path,num_classes,train=True,size=512,l=(16,16)):
        """
        :param path:
        """
        self.l=l
        self.num_classes=num_classes
        if train:
            self.images_path= os.path.join(path,"train","images")
            self.labels_path = os.path.join(path, "train", "labels")
        else:
            self.images_path = os.path.join(path, "test", "images")
            self.labels_path = os.path.join(path, "test", "labels")
        images_files = os.listdir(self.images_path)
        labels_files = os.listdir(self.labels_path)
        images_files = list(map(lambda x:(x,x.split(".")[0],len(x.split("."))),images_files))
        labels_files = list(map(lambda x:(x,x.split(".")[0],len(x.split("."))),labels_files))
        images_files=[i[:-1] for i in images_files if i[-1]==2]
        labels_files = [i[:-1] for i in labels_files if i[-1] == 2]
        self.files=[]
        for i in images_files:
            for j in labels_files:
                if i[1]==j[1]:
                    self.files.append((i[0],j[0]))
                    break
        self.transform=transforms.Compose([
            transforms.Resize((size,size)),
            transforms.ToTensor()
        ])
    def __getitem__(self, item):
        image=Image.open(os.path.join(self.images_path,self.files[item][0]))
        image=self.transform(image)
        label=self.read_txt(os.path.join(self.labels_path,self.files[item][1]))
        label=utils.get_label(label,num_classes=self.num_classes,l=self.l)
        return image,label
    def __len__(self):
        return len(self.files)
    def read_txt(self,path):
        """
        class,x,y,h,w
        """
        all=[]
        with open(path,"r",encoding="utf-8") as F:
            for line in F.readlines():
                a=line.split(" ")
                class_=[int(a[0])]
                x_y_h_w=list(map(float,a[1:]))
                all.append(class_+x_y_h_w)
        return all

if __name__ == '__main__':
    pass