File size: 6,242 Bytes
25986db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | import os
import os.path
import torch
import numpy as np
import pandas
import csv
from collections import OrderedDict
from .base_video_dataset import BaseVideoDataset
from lib.train.data import jpeg4py_loader_w_failsafe
from lib.train.admin import env_settings
from lib.train.dataset.depth_utils import get_x_frame
class DepthTrack(BaseVideoDataset):
""" DepthTrack dataset.
"""
def __init__(self, root=None, dtype='rgbcolormap', split='train', image_loader=jpeg4py_loader_w_failsafe,
multi_modal_vision=False, multi_modal_language=False): # vid_ids=None, split=None, data_fraction=None
"""
args:
image_loader (jpeg4py_loader) - The function to read the images. jpeg4py (https://github.com/ajkxyz/jpeg4py)
is used by default.
vid_ids - List containing the ids of the videos (1 - 20) used for training. If vid_ids = [1, 3, 5], then the
videos with subscripts -1, -3, and -5 from each class will be used for training.
# split - If split='train', the official train split (protocol-II) is used for training. Note: Only one of
# vid_ids or split option can be used at a time.
# data_fraction - Fraction of dataset to be used. The complete dataset is used by default
root - path to the lasot depth dataset.
dtype - colormap or depth, colormap + depth
if colormap, it returns the colormap by cv2,
if depth, it returns [depth, depth, depth]
"""
root = env_settings().depthtrack_dir if root is None else root
super().__init__('DepthTrack', root, image_loader)
self.dtype = dtype # colormap or depth
self.split = split
self.sequence_list = self._build_sequence_list()
self.seq_per_class, self.class_list = self._build_class_list()
self.class_list.sort()
self.class_to_id = {cls_name: cls_id for cls_id, cls_name in enumerate(self.class_list)}
self.multi_modal_vision = multi_modal_vision
self.multi_modal_language = multi_modal_language
def _build_sequence_list(self):
ltr_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
file_path = os.path.join(ltr_path, 'data_specs', 'depthtrack_%s.txt'%self.split)
# sequence_list = pandas.read_csv(file_path, header=None, squeeze=True).values.tolist()
sequence_list = pandas.read_csv(file_path, header=None).squeeze("columns").values.tolist()
return sequence_list
def _build_class_list(self):
seq_per_class = {}
class_list = []
for seq_id, seq_name in enumerate(self.sequence_list):
class_name = seq_name.split('_')[0]
if class_name not in class_list:
class_list.append(class_name)
if class_name in seq_per_class:
seq_per_class[class_name].append(seq_id)
else:
seq_per_class[class_name] = [seq_id]
return seq_per_class, class_list
def get_name(self):
return 'depthtrack'
def has_class_info(self):
return True
def has_occlusion_info(self):
return True
def get_num_sequences(self):
return len(self.sequence_list)
def get_num_classes(self):
return len(self.class_list)
def get_sequences_in_class(self, class_name):
return self.seq_per_class[class_name]
def _read_bb_anno(self, seq_path):
bb_anno_file = os.path.join(seq_path, "groundtruth.txt")
gt = pandas.read_csv(bb_anno_file, delimiter=',', header=None, dtype=np.float32, na_filter=True, low_memory=False).values
return torch.tensor(gt)
def _get_sequence_path(self, seq_id):
seq_name = self.sequence_list[seq_id]
return os.path.join(self.root, seq_name,seq_name)
def get_sequence_info(self, seq_id):
seq_path = self._get_sequence_path(seq_id)
bbox = self._read_bb_anno(seq_path) # xywh just one kind label
'''
if the box is too small, it will be ignored
'''
# valid = (bbox[:, 2] > 0) & (bbox[:, 3] > 0)
valid = (bbox[:, 2] > 10.0) & (bbox[:, 3] > 10.0)
visible = valid.clone().byte()
return {'bbox': bbox, 'valid': valid, 'visible': visible}
def _get_frame_path(self, seq_path, frame_id):
'''
return depth image path
'''
return os.path.join(seq_path, 'color', '{:08}.jpg'.format(frame_id+1)) , os.path.join(seq_path, 'depth', '{:08}.png'.format(frame_id+1)) # frames start from 1
def _get_frame(self, seq_path, frame_id):
'''
Return :
- colormap from depth image
- 3xD = [depth, depth, depth], 255
- rgbcolormap
- rgb3d
- color
- raw_depth
'''
color_path, depth_path = self._get_frame_path(seq_path, frame_id)
img = get_x_frame(color_path, depth_path, dtype=self.dtype, depth_clip=True)
return img
def _get_class(self, seq_path):
# raw_class = seq_path.split('/')[-2]
# return raw_class
return self.split
def get_class_name(self, seq_id):
depth_path = self._get_sequence_path(seq_id)
obj_class = self._get_class(depth_path)
return obj_class
def get_frames(self, seq_id, frame_ids, anno=None):
seq_path = self._get_sequence_path(seq_id)
obj_class = self._get_class(seq_path)
if anno is None:
anno = self.get_sequence_info(seq_id)
anno_frames = {}
for key, value in anno.items():
anno_frames[key] = [value[f_id, ...].clone() for ii, f_id in enumerate(frame_ids)]
frame_list = [self._get_frame(seq_path, f_id) for ii, f_id in enumerate(frame_ids)]
object_meta = OrderedDict({'object_class_name': obj_class,
'motion_class': None,
'major_class': None,
'root_class': None,
'motion_adverb': None})
return frame_list, anno_frames, object_meta
|