Data loading for the UCSD pedestrian counting dataset

UCSD_Ped_Counting_Dataset.py

class UCSD(Dataset):
    """UCSD pedestrian counting data."""

    def __init__(self, data_dir, annotation_dir, transform=None):

        self.file_list = []
        self.file_cnts = []
        
        files = glob.glob(
            os.path.join(
                annotation_dir, 
                '*count_roi_mainwalkway.mat'))
        
        for f in files:
            tmp = loadmat(f)
            
            l_count = tmp['count'][0][0].ravel()
            r_count = tmp['count'][0][1].ravel()
            t_count = l_count + r_count
            
            [self.file_cnts.append(c) for c in t_count]
            
            file_parts = os.path.basename(f).split('_')
            seq_id = "_".join(file_parts[0:3])
            
            for i in np.arange(len(t_count)):
                
                self.file_list.append(
                    os.path.join(
                        data_dir,
                        seq_id + ".y",
                        "{}_f{:03d}.png".format(seq_id,i+1)))
            
            self.transform = transform

    def __len__(self):
        return len(self.file_list)
        
    def __getitem__(self, idx):
        img_name = os.path.join(self.file_list[idx])
        img = Image.open(img_name)
        img = img.resize((128,128))

        if self.transform:
            img = self.transform(img)
        return img, self.file_cnts[idx]

Thanks for the script @rkwitt.

I have made a few changes in my fork like import statements. You can add them.
So, others can use it on the go.

Again, Thanks very much.