以STARK中的sampler.py为例:
一、代码整体结构
代码整体结构如下,分为以下阶段:
# 1. 任选一个数据集,如Lasot; # 2. 从该数据集中任选一个视频序列,如airplane-5; # 3. 从该视频序列中任选一个base frame; # 4. 分别从[base_frame_id - max_gap, base_frame_id]和(base_frame_id, base_frame_id + max_gap]中采样一系列train frames和test frames, 需要注意的是只有visible的帧才会被采样; # 5. 如果没有找到满足4中的这些帧,逐渐增加max_gap直到这些帧都找到; # 6. 找到这些帧后对其进行预处理。
class TrackingSampler(torch.utils.data.Dataset):
# 采样帧来形成一个一个batch。
# 采样过程分为以下几个阶段:
# 1. 任选一个数据集,如Lasot;
# 2. 从该数据集中任选一个视频序列,如airplane-5;
# 3. 从该视频序列中任选一个base frame;
# 4. 分别从[base_frame_id - max_gap, base_frame_id]和(base_frame_id, base_frame_id + max_gap]中采样一系列train frames和test frames, 需要注意的是只有visible的帧才会被采样;
# 5. 如果没有找到满足4中的这些帧,逐渐增加max_gap直到这些帧都找到;
# 6. 找到这些帧后对其进行预处理。
def __init__(self, datasets, p_datasets, samples_per_epoch, max_gap,
num_search_frames, num_template_frames=1, processing=no_processing, frame_sample_mode='causal',
train_cls=False, pos_prob=0.5):
def __len__(self):
def _sample_visible_ids(self, visible, num_ids=1, min_id=None, max_id=None,
allow_invisible=False, force_invisible=False):
# 在min_id和max_id之间采样num_ids个visible的帧。
def __getitem__(self, index):
if self.train_cls:
return self.getitem_cls() # 增加时序信息时使用这个,即score head的训练,用于更新模板
else:
return self.getitem() # 无时序信息时用这个
def getitem(self):
# 下一模块详细阐述
def getitem_cls(self):
# 下一模块详细阐述
def get_center_box(self, H, W, ratio=1/8):
def sample_seq_from_dataset(self, dataset, is_video_dataset):
# 采样一个有足够visible frames的视频序列
def get_one_search(self):
def get_frame_ids_trident(self, visible):
def get_frame_ids_stark(self, visible, valid):
二、 def getitem_cls(self)
sampler.py中getitem_cls(self)函数的解释如下:
# 1. 任选一个数据集(以概率p_datasets选取) # 2. 从该数据集中采样有足够visible frames的视频序列 # 3. 在该视频序列中确定template_frame_ids和search_frame_ids(要求,visible且间距不超过max_gap) # 4. 根据template_frame_ids获取templates的图像和bbox,valid,visible,class等信息 # 5. 获取正样本/负样本的search用于分类。正样本直接将第3步得到的search_frame_ids作为search,此时label=1; 负样本则重新任选数据集、视频序列、visible frame作为search,此时label=0。 # 6. 数据增强
def getitem_cls(self):
# 获取用于分类的数据。在STARK中即用于score head的训练
valid = False
label = None
while not valid:
# 1. 任选一个数据集(以概率p_datasets选取)
dataset = random.choices(self.datasets, self.p_datasets)[0] # lasot
is_video_dataset = dataset.is_video_sequence() # Ture
# 2. 从该数据集中采样有足够visible frames的视频序列
seq_id, visible, seq_info_dict = self.sample_seq_from_dataset(dataset, is_video_dataset)
# seq_id=4, visible全为1的tensor,seq_info_dict字典是该序列所有帧的gt信息{'bbox','valid','visible'}
# 3. 在该视频序列中确定template_frame_ids和search_frame_ids(要求,visible且间距不超过max_gap)
if is_video_dataset:
if self.frame_sample_mode in ["trident", "trident_pro"]:
# 这里self.frame_sample_mode是"trident_pro"
template_frame_ids, search_frame_ids = self.get_frame_ids_trident(visible)
# 采样得到间距不超过max_gap的三个visible frames,如template_frame_ids=[4159,4078],search_frame_ids=[4250]
elif self.frame_sample_mode == "stark":
template_frame_ids, search_frame_ids = self.get_frame_ids_stark(visible, seq_info_dict["valid"])
else:
raise ValueError("illegal frame sample mode")
else:
# In case of image dataset, just repeat the image to generate synthetic video
template_frame_ids = [1] * self.num_template_frames
search_frame_ids = [1] * self.num_search_frames
# 根据template_frame_ids和search_frame_ids找到其images和bbox
try:
# "try" is used to handle trackingnet data failure
# 4. 根据template_frame_ids获取templates的图像和bbox,valid,visible,class等信息
template_frames, template_anno, meta_obj_train = dataset.get_frames(seq_id, template_frame_ids,
seq_info_dict)
H, W, _ = template_frames[0].shape
template_masks = template_anno['mask'] if 'mask' in template_anno else [torch.zeros((H, W))] * self.num_template_frames # 因为lasot中没有mask注释,所以template_masks全为0
# 5. 获取正样本/负样本的search用于分类。正样本直接将第3步得到的search_frame_ids作为search,此时label=1; 负样本则重新任选数据集、视频序列、visible frame作为search,此时label=0。
# positive samples
if random.random() < self.pos_prob:
label = torch.ones(1,)
search_frames, search_anno, meta_obj_test = dataset.get_frames(seq_id, search_frame_ids, seq_info_dict)
search_masks = search_anno['mask'] if 'mask' in search_anno else [torch.zeros(
(H, W))] * self.num_search_frames
# negative samples
else:
label = torch.zeros(1,)
if is_video_dataset:
search_frame_ids = self._sample_visible_ids(visible, num_ids=1, force_invisible=True)
if search_frame_ids is None:
# 任选一个数据集、视频序列、任选一个visible frame作为search
search_frames, search_anno, meta_obj_test = self.get_one_search()
else:
search_frames, search_anno, meta_obj_test = dataset.get_frames(seq_id, search_frame_ids,
seq_info_dict)
search_anno["bbox"] = [self.get_center_box(H, W)]
else:
search_frames, search_anno, meta_obj_test = self.get_one_search()
H, W, _ = search_frames[0].shape
search_masks = search_anno['mask'] if 'mask' in search_anno else [torch.zeros(
(H, W))] * self.num_search_frames
data = TensorDict({'template_images': template_frames,
'template_anno': template_anno['bbox'],
'template_masks': template_masks,
'search_images': search_frames,
'search_anno': search_anno['bbox'],
'search_masks': search_masks,
'dataset': dataset.get_name(),
'test_class': meta_obj_test.get('object_class_name')})
# 数据增强
data = self.processing(data)
# add classification label
data["label"] = label # 0为负样本
# check whether data is valid
valid = data['valid']
except:
valid = False
return data