import json
import os
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import numpy as np
import random
import pdb
from PIL import Image
def _read_labels(json_file, imus, sequence_length):
"""Returns a list of all frames, and a list of where each data point (whose
length is sequence_length) in the list of frames."""
with open(json_file, 'r') as fp:
dataset_meta = json.load(fp)
frames = []
idx_to_fid = []
centroids = {
'absolute_centroids': torch.Tensor(dataset_meta['absolute_centroids']),
'difference_centroids':
torch.Tensor(dataset_meta['difference_centroids']),
}
for clip in dataset_meta['clips']:
pdb.set_trace()
frame_clips = [{
'cur_frame':
frame_meta['filename'],
'prev_frame':
frame_meta['prev-frame'],
'labels':
torch.LongTensor(
[frame_meta['imu-diff-clusters'][imu] for imu in imus]),
'diffs':
torch.FloatTensor(
[frame_meta['imu-diff-values'][imu] for imu in imus]),
'absolute_cur_imus':
torch.FloatTensor(
[frame_meta['absolute_cur_imus'][imu] for imu in imus]),
'absolute_prev_imus':
torch.FloatTensor(
[frame_meta['absolute_prev_imus'][imu] for imu in imus]),
'imu_diff_values':
torch.FloatTensor(
[frame_meta['imu-diff-values'][imu] for imu in imus])
} for frame_meta in clip['frames']]
for i in range(len(frame_clips) - sequence_length + 1):
idx_to_fid.append(i + len(frames))
frames += frame_clips
return frames, idx_to_fid, centroids
def _category_weights():
category_sizes = torch.Tensor(
[[131., 199., 177., 157., 3446., 1689., 14838.,
186.], [379., 366., 1705., 1297., 9746., 873., 1475.,
4982.], [232., 257., 241., 3422., 126., 11225., 5105., 215.],
[137., 115., 142., 3192., 3066., 10036., 3983.,
152.], [781., 594., 1183., 4753., 664., 9038., 1394.,
2416.], [140., 214., 150., 4644., 11075., 1974., 182., 2444.],
[169., 133., 150., 104., 2090., 12701., 2157., 3319.]])
weight = category_sizes.sum(1, keepdim=True) / category_sizes
return weight
class DogClipRegrDataset(data.Dataset):
CLASS_WEIGHTS = _category_weights()
def __init__(self, args, train=True):
root_dir = args.data
if train or args.read_feature_and_image:
json_file = os.path.join(root_dir, 'train.json')
elif args.use_test_for_val:
json_file = os.path.join(root_dir, 'test.json')
else:
json_file = os.path.join(root_dir, 'val.json')
self.num_classes = args.num_classes
self.sequence_length = args.sequence_length
self.experiment_type = args.experiment_type
self.regression = args.regression
self.args = args
self.read_features = args.read_features
self.frames_metadata, self.idx_to_fid, self.centroids = _read_labels(
json_file, args.imus, args.sequence_length)
self.root_dir = root_dir
self.features_dir = args.features_dir
self.transform = transforms.Compose([
transforms.Scale((args.image_size, args.image_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
def get_relative_centroids(self):
return self.centroids['difference_centroids']
def __len__(self):
return len(self.idx_to_fid)
def load_and_resize(self, img_name):
with open(img_name, 'rb') as fp:
image = Image.open(fp).convert('RGB')
return self.transform(image)
def get_relpath(self, idx):
fid = self.idx_to_fid[idx]
return self.frames_metadata[fid]['cur_frame']
def __getitem__(self, idx):
current_images_files = []
fid = self.idx_to_fid[idx]
if self.read_features:
features = []
for i in range(self.sequence_length):
feature_path = os.path.join(
self.features_dir,
self.frames_metadata[fid + i]['cur_frame'] + '.pytar')
features.append(torch.load(feature_path))
input = torch.stack(features)
if not self.read_features:
cur_frame = self.load_and_resize(
os.path.join(self.root_dir, 'images',
self.frames_metadata[fid]['cur_frame']))
prev_frame = self.load_and_resize(
os.path.join(self.root_dir, 'images',
self.frames_metadata[fid]['prev_frame']))
input = torch.cat([cur_frame, prev_frame], 0)
labels = []
absolute_cur_imus = []
absolute_prev_imus = []
for i in range(self.sequence_length):
if self.regression:
labels.append(self.frames_metadata[fid + i]['diffs'])
else:
labels.append(self.frames_metadata[fid + i]['labels'])
absolute_cur_imus.append(
self.frames_metadata[fid + i]['absolute_cur_imus'])
absolute_prev_imus.append(
self.frames_metadata[fid + i]['absolute_prev_imus'])
current_images_files.append(
self.frames_metadata[fid + i]['cur_frame'])
labels = torch.stack(labels)
absolute_cur_imus = torch.stack(absolute_cur_imus)
absolute_prev_imus = torch.stack(absolute_prev_imus)
if self.args.read_feature_and_image:
features = []
for i in range(self.sequence_length):
image = self.load_and_resize(
os.path.join(self.root_dir, 'images',
self.frames_metadata[fid + i]['cur_frame']))
features.append(image)
labels = torch.stack(features)
pdb.set_trace()
return (input, labels, absolute_prev_imus, absolute_cur_imus,
current_images_files)
