Data loading for the UCSD pedestrian counting(http://www.svcl.ucsd.edu/projects/peoplecnt/) dataset for images with their respective count from matlab scripts provided along.

UCSD_count_preprocess.py

import glob
import os
import scipy.io as sio
import numpy as np
import numpy as np
from PIL import *

class UCSD():
    """UCSD pedestrian counting data."""
    
    def __init__(self, data_dir, annotation_dir, transform=None):
        
        # stores name of files and their respective counts in two different lists.
        self.file_list = []
        self.file_cnts = []
        
        files = glob.glob(
            os.path.join(
                annotation_dir, 
                '*count_roi_mainwalkway.mat'))
        
        for f in files:
            tmp = sio.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]
        

# For loading the dataset into an object.
X = UCSD('ucsdpeds/', 'vidf-cvpr/')

dict_imgs = dict(zip(X.file_list, X.file_cnts))