class FeatureDataset(Dataset):
def __init__(self, phase):
if phase == "train":
self.tsv_path = TRAIN_PATH
elif phase == 'val':
self.tsv_path = VAL_PATH
else:
self.tsv_path = TEST_PATH
self.tsv = TSVFile(self.tsv_path)
def __getitem__(self, idx):
row = self.tsv.seek(idx)
# ginfo = row[0]
glabel = row[0]
actions = row[1][1:-1].split(',')
actions = map(int, actions)
features = json.loads(row[2])
person_features = []
for pid in features:
person_feature = np.array(features[pid])
person_features.append(person_feature)
return person_features, actions, glabel
# return person_features, actions, glabel, ginfo
def __len__(self):
return self.tsv.num_rows()
class TSVDetection(data.Dataset):
def __init__(self, tsv_file, labelmap, transform=None):
from tsv_io import TSVFile
from qd_common import load_list_file
self._tsv = TSVFile(tsv_file)
self._label_to_idx = {
l: i
for i, l in enumerate(load_list_file(labelmap))
}
self.transform = transform
self.ids = None
def __getitem__(self, index):
key, im, gt, h, w = self.pull_item(index)
return im, gt
def get_ids(self):
if self.ids is None:
self.ids = [key for key, _, __ in self._tsv]
return self.ids
def __len__(self):
return self._tsv.num_rows()
def pull_item(self, index):
from qd_common import img_from_base64
key, label_str, img_str = self._tsv.seek(index)
img = img_from_base64(img_str)
height, width, channels = img.shape
rects = json.loads(label_str)
rects = [r for r in rects if 'diff' not in r or not r['diff']]
for r_info in rects:
r = r_info['rect']
r[0], r[2] = r[0] / width, r[2] / width
r[1], r[3] = r[1] / height, r[3] / height
idx = self._label_to_idx[r_info['class']]
r.append(idx)
target = [r['rect'] for r in rects]
if self.transform is not None:
target = np.array(target)
img, boxes, labels = self.transform(img, target[:, :4], target[:,
4])
# to rgb
img = img[:, :, (2, 1, 0)]
# img = img.transpose(2, 0, 1)
target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
return key, torch.from_numpy(img).permute(2, 0,
1), target, height, width
class VolleyballDataset(Dataset):
def __init__(self, phase):
self.phase = phase
self.image_path = "../dataset/Volleyball/videos/"
if phase == "trainval":
self.tsv_path = TRAINVAL_PATH
else:
self.tsv_path = TEST_PATH
self.tsv = TSVFile(self.tsv_path)
def __getitem__(self, idx):
num_boxes = 12
row = self.tsv.seek(idx)
json_dict = json.loads(row[0])
set_idx = json_dict['video']
target_fidx = int(json_dict['target'])
frame_idxes = int(json_dict['frame'])
group_activities = np.array(json_dict['glabel'])
group_info = np.array((set_idx, target_fidx, frame_idxes))
person_actions = np.array(json_dict['plabel'])
person_num = person_actions.shape[0]
if person_num != num_boxes:
person_actions = np.hstack(
[person_actions, person_actions[:num_boxes - person_num]])
person_actions = torch.tensor(person_actions)
bboxes = np.array(json_dict['bboxes'])
bboxes_total = np.concatenate(
[bboxes, bboxes[:num_boxes - person_num]], axis=0)
img_tensor = []
img_path = self.image_path + '{}/{}/{}.jpg'.format(
set_idx, target_fidx, frame_idxes)
frame_img = Image.open(img_path)
if set_idx in HIGH_RESOLUTION:
img_tensor.append(resized_transform(frame_img))
else:
img_tensor.append(transform(frame_img))
# (B, 3, 780, 1280)
return torch.stack(
img_tensor
), bboxes_total, person_actions, group_activities, group_info
def __len__(self):
return self.tsv.num_rows()
class VolleyballDataset(Dataset):
def __init__(self, phase):
if phase == "train":
self.tsv_path = TRAIN_PATH
elif phase == 'val':
self.tsv_path = VAL_PATH
else:
self.tsv_path = TEST_PATH
self.tsv = TSVFile(self.tsv_path)
def __getitem__(self, idx):
row = self.tsv.seek(idx)
set_id = row[0]
tar_frame = row[1]
gact = row[2]
group_info = (set_id, tar_frame, gact)
group_dict = json.loads(row[3])
person_num = len(group_dict)
person_actions = []
person_pixels = []
# person_pixels = {pid: [] for pid in range(person_num)}
int_pid = 0
for pid in group_dict:
person_actions.append(group_dict[pid][0])
for i in range(1, NUM_FRAMES + 1):
img_np = np.array(group_dict[pid][i])
image = Image.fromarray(img_np.astype('uint8'), 'RGB')
img_tensor = trans(image)
person_pixels.append(img_tensor)
# person_pixels[int_pid].append(img_tensor)
int_pid += 1
return person_pixels, person_actions, group_info
def __len__(self):
return self.tsv.num_rows()
class VolleyballDataset(Dataset):
def __init__(self, phase, frame_mode):
self.phase = phase
if frame_mode == 'target-frame':
TRAINVAL_PATH = './data/feature drop-0.2 trainval.tsv'
TEST_PATH = './data/feature drop-0.2 test.tsv'
elif frame_mode == 'frames':
TRAINVAL_PATH = './data/feature frames drop-0.2 trainval.tsv'
TEST_PATH = './data/feature frames drop-0.2 test.tsv'
else:
print('wrong frame mode')
sys.exit(0)
if phase == 'trainval':
self.tsv_path = TRAINVAL_PATH
elif phase == 'test':
self.tsv_path = TEST_PATH
else:
print('wrong dataset phase')
sys.exit(0)
self.tsv = TSVFile(self.tsv_path)
def __getitem__(self, idx):
row = self.tsv.seek(idx)
json_dict = json.loads(row[0])
group_info = torch.tensor(json_dict['group_info'])
actions = torch.tensor(
json_dict['actions']).view(-1).repeat(NUM_FRAMES)
# actions = torch.tensor(json_dict['actions']).repeat(NUM_FRAMES)
activities = torch.tensor(json_dict['activities']).repeat(NUM_FRAMES)
features = torch.tensor(json_dict['featuers'])
return group_info, actions, activities, features
def __len__(self):
return self.tsv.num_rows()
class VolleyballDataset(Dataset):
def __init__(self, phase):
self.phase = phase
if phase == "trainval":
self.tsv_path = TRAINVAL_PATH
else:
self.tsv_path = TEST_PATH
self.tsv = TSVFile(self.tsv_path)
def __getitem__(self, idx):
row = self.tsv.seek(idx)
json_dict = json.loads(row[0])
group_info = torch.tensor(json_dict['group_info'])
actions = torch.tensor(json_dict['actions']).repeat(NUM_FRAMES)
activities = np.array(json_dict['activities']).repeat(NUM_FRAMES)
features = torch.tensor(json_dict['featuers'])
return group_info, actions, activities, features
def __len__(self):
return self.tsv.num_rows()
class VolleyballDataset(Dataset):
def __init__(self, phase):
self.phase = phase
if phase == "trainval":
self.tsv_path = TRAINVAL_PATH
else:
self.tsv_path = TEST_PATH
self.tsv = TSVFile(self.tsv_path)
def __getitem__(self, idx):
num_boxes = 12
row = self.tsv.seek(idx)
json_dict = json.loads(row[0])
set_id = json_dict['video']
tar_frame = json_dict['tf']
group_activities = json_dict['glabel']
group_info = np.array((set_id, tar_frame, group_activities))
person_actions = np.array(json_dict['plabel'])
person_num = person_actions.shape[0]
bboxes = np.array(json_dict['bboxes'])
if person_num != num_boxes:
bboxes = np.vstack([bboxes, bboxes[:num_boxes - len(bboxes[-1])]])
person_actions = np.hstack(
[person_actions, person_actions[:num_boxes - person_num]])
img_np = np.array(json_dict['pixels'])
image = Image.fromarray(img_np.astype('uint8'), 'RGB')
img_tensor = test_transform(image)
return img_tensor, bboxes, person_actions, np.array(
group_activities), group_info
def __len__(self):
return self.tsv.num_rows()