Skip to content

Instantly share code, notes, and snippets.

@rkwitt
Created September 25, 2017 09:42
Show Gist options
  • Save rkwitt/0c781d1c3cc9ae5db8cdf35fa30648df to your computer and use it in GitHub Desktop.
Save rkwitt/0c781d1c3cc9ae5db8cdf35fa30648df to your computer and use it in GitHub Desktop.
Data loading for the UCSD pedestrian counting dataset
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]
@ashishrana160796
Copy link

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment