def test_flowwrite():
flow = np.random.rand(100, 100, 2).astype(np.float32)
# write to a .flo file
_, filename = tempfile.mkstemp()
mmcv.flowwrite(flow, filename)
flow_from_file = mmcv.flowread(filename)
assert_array_equal(flow, flow_from_file)
os.remove(filename)
# write to two .jpg files
tmp_filename = osp.join(tempfile.gettempdir(), 'mmcv_test_flow.jpg')
for concat_axis in range(2):
mmcv.flowwrite(flow,
tmp_filename,
quantize=True,
concat_axis=concat_axis)
shape = (200, 100) if concat_axis == 0 else (100, 200)
assert osp.isfile(tmp_filename)
assert mmcv.imread(tmp_filename, flag='unchanged').shape == shape
os.remove(tmp_filename)
# test exceptions
with pytest.raises(AssertionError):
mmcv.flowwrite(flow, tmp_filename, quantize=True, concat_axis=2)
def prepare_clip_train(self, idx):
vid, fid = self.sample_list[idx]
video = self.video_infos[vid]
frames = video['frames']
quantize = video['quantize']
if self.padding:
pass
else:
frames = frames[fid:fid + self.seq_len]
# img_names = [x['filename'] for x in frames]
flows = [
mmcv.flowread(osp.join(self.img_prefix, x['flow_name']),
quantize=quantize) for x in frames
]
inv_flows = [
mmcv.flowread(osp.join(self.img_prefix, x['inv_flow_name']),
quantize=quantize) for x in frames
]
img_results = [self.prepare_img_train(frame) for frame in frames]
# results encode the transformation info of img
# augment the flow/inv_flow accordingly, pack it into results
# resize => flip => pad
for i, x in enumerate(img_results):
if x is None:
return None
aug_meta = img_results[0]['img_meta'].data
resize = aug_meta['img_shape'][:2]
pad = aug_meta['pad_shape'][:2]
flip = aug_meta['flip']
flow_aug = FlowAug(resize, pad, flip)
flows = [torch.tensor(x).permute(2, 0, 1) for x in flows]
inv_flows = [torch.tensor(x).permute(2, 0, 1) for x in inv_flows]
flows = torch.stack(flows)
inv_flows = torch.stack(inv_flows)
flows, inv_flows = flow_aug(flows, inv_flows)
for idx, img_result in enumerate(img_results):
# if not isinstance(img_result, dict):
# print("Debug")
# print(img_result)
# print(frames[idx])
# print(frames)
img_result['flow'] = flows[idx]
img_result['inv_flow'] = inv_flows[idx]
# embed()
return img_results
def get_flow_blob(target_scale, target_max_size, vid_frm_idx):
vid, str_frm_idx = vid_frm_idx.split('/')
cur_frm_idx = int(str_frm_idx)
flows = []
shared_flow_scales = None
if cfg.A2D.LOAD_FLOW:
for f_idx in range(cur_frm_idx - int(cfg.A2D.SEGMENT_LENGTH / 2),
cur_frm_idx + int(cfg.A2D.SEGMENT_LENGTH / 2)):
flow_path = os.path.join(cfg.A2D.FLOW_ROOT, vid,
'%05d' % f_idx + '.flo')
assert os.path.exists(flow_path)
flow = mmcv.flowread(flow_path)
flows_with_diff_scales, computed_flow_scales = prep_flow_for_blob(
flow,
cfg.A2D.FLOW_MAX_MAG, [target_scale],
target_max_size,
clip_mag=cfg.A2D.CLIP_FLOW_MAG)
three_channel_flow = flow_to_flow_img(flows_with_diff_scales[0])
flows.append(three_channel_flow)
if shared_flow_scales is not None:
assert computed_flow_scales[0] == shared_flow_scales
else:
shared_flow_scales = computed_flow_scales[0]
else:
prev_frame = None
for f_idx in range(cur_frm_idx - int(cfg.A2D.SEGMENT_LENGTH / 2),
cur_frm_idx + int(cfg.A2D.SEGMENT_LENGTH / 2) + 1):
frame_fpath = os.path.join(cfg.A2D.ROOT, vid,
'%05d' % f_idx + '.png')
assert os.path.exists(frame_fpath)
cur_frame = cv2.imread(frame_fpath)
cur_frame = cv2.cvtColor(cur_frame, cv2.COLOR_BGR2GRAY)
if prev_frame is not None:
flow = cv2.calcOpticalFlowFarneback(prev_frame, cur_frame,
None, 0.5, 3, 15, 3, 5,
1.2, 0)
flows_with_diff_scales, computed_flow_scales = prep_flow_for_blob(
flow,
cfg.A2D.FLOW_MAX_MAG, [target_scale],
target_max_size,
clip_mag=cfg.A2D.CLIP_FLOW_MAG)
three_channel_flow = flow_to_flow_img(
flows_with_diff_scales[0])
flows.append(three_channel_flow)
if shared_flow_scales is not None:
assert computed_flow_scales[0] == shared_flow_scales
else:
shared_flow_scales = computed_flow_scales[0]
prev_frame = cur_frame
blob = flow_list_to_blob([flows])
height, width = blob.shape[2], blob.shape[3]
im_info = np.hstack((height, width, shared_flow_scales))[np.newaxis, :]
# return blob, shared_flow_scales, im_info.astype(np.float32)
return blob
def test_flow_from_bytes():
data_dir = osp.join(osp.dirname(__file__), '../data')
flow_shape = (60, 80, 2)
flow_file = osp.join(data_dir, 'optflow.flo')
# read .flo file
flow_fromfile = mmcv.flowread(flow_file)
with open(flow_file, 'rb') as f:
flow_bytes = f.read()
flow_frombytes = mmcv.flow_from_bytes(flow_bytes)
assert flow_frombytes.shape == flow_shape
assert np.all(flow_frombytes == flow_fromfile)
def read_flow_array(path: str, return_info, to_normalize=False):
"""
:param path: 支持flo(mmcv)和png(opencv)数据
:param return_info: 对于flo数据，支持三种选择(flow,direction,magnitude)，对于png数据直接读取，返回键为flow
:param to_normalize: 仅对flo数据的magnitude数据有效
:return: 0~1
"""
if path.endswith('.flo'):
flow_array = mmcv.flowread(path)
split_flow = _flow_to_direction_and_magnitude(flow_array)
if not isinstance(return_info, (tuple, list)):
return_info = [return_info]
return_array = dict()
for k in return_info:
data_array = split_flow[k]
if k == 'magnitude' and to_normalize:
data_array = (data_array - data_array.min()) / \
(data_array.max() - data_array.min())
return_array[k] = data_array
else:
return_array = dict(flow=read_color_array(path))
return return_array
numpy.array(PIL.Image.open(img_second))[:, :, ::-1].transpose(
2, 0, 1).astype(numpy.float32) * (1.0 / 255.0))
tensorOutput = estimate(tensorFirst, tensorSecond)
print(arguments_strFolOut, directory_name)
outfile = os.path.join(arguments_strFolOut, '%06d.flo' % (i))
objectOutput = open(outfile, 'wb')
numpy.array([80, 73, 69, 72], numpy.uint8).tofile(objectOutput)
numpy.array([tensorOutput.size(2),
tensorOutput.size(1)],
numpy.int32).tofile(objectOutput)
numpy.array(tensorOutput.numpy().transpose(1, 2, 0),
numpy.float32).tofile(objectOutput)
objectOutput.close()
if arguments_strVisFlow:
flow = mmcv.flow2rgb(mmcv.flowread(outfile))
io.imsave(outfile[0:-3] + 'png', flow)
else:
tensorFirst = torch.FloatTensor(
numpy.array(
PIL.Image.open(arguments_strFirst))[:, :, ::-1].transpose(
2, 0, 1).astype(numpy.float32) * (1.0 / 255.0))
tensorSecond = torch.FloatTensor(
numpy.array(
PIL.Image.open(arguments_strSecond))[:, :, ::-1].transpose(
2, 0, 1).astype(numpy.float32) * (1.0 / 255.0))
tensorOutput = estimate(tensorFirst, tensorSecond)
objectOutput = open(arguments_strOut, 'wb')
# patch = mmcv.imcrop(img1, bboxes1)
#
# # crop two regions (10, 10, 100, 120) and (0, 0, 50, 50)
# bboxes2 = np.array([[10, 10, 100, 120], [0, 0, 50, 50]])
# patches1 = mmcv.imcrop(img1, bboxes2)
#
# # crop two regions, and rescale the patches by 1.2x
# patches2 = mmcv.imcrop(img1, bboxes2)
#
# mmcv.imshow(patches2[1])
# img = mmcv.impad(img2, shape=[600, 1000], pad_val=200)
#
# mmcv.imshow(img)
#video = mmcv.VideoReader('/Users/user/Desktop/datatset/P03_coffee.avi')
# obtain basic informationConcise
# print(len(video))
# print(video.width, video.height, video.resolution, video.fps)
#
#
# # read some frames
# img = video[100]
# mmcv.imshow(img)
#video.cvt2frames('./frames', start=500)
flow = mmcv.flowread('test.flo')
mmcv.flowshow(flow)
def test_flowread():
flow_shape = (60, 80, 2)
# read .flo file
flow = mmcv.flowread(osp.join(osp.dirname(__file__), 'data/optflow.flo'))
assert flow.shape == flow_shape
# pseudo read
flow_same = mmcv.flowread(flow)
assert_array_equal(flow, flow_same)
# read quantized flow concatenated vertically
flow = mmcv.flowread(osp.join(osp.dirname(__file__),
'data/optflow_concat0.jpg'),
quantize=True,
denorm=True)
assert flow.shape == flow_shape
# read quantized flow concatenated horizontally
flow = mmcv.flowread(osp.join(osp.dirname(__file__),
'data/optflow_concat1.jpg'),
quantize=True,
concat_axis=1,
denorm=True)
assert flow.shape == flow_shape
# test exceptions
notflow_file = osp.join(osp.dirname(__file__), 'data/color.jpg')
with pytest.raises(TypeError):
mmcv.flowread(1)
with pytest.raises(IOError):
mmcv.flowread(notflow_file)
with pytest.raises(IOError):
mmcv.flowread(notflow_file, quantize=True)
with pytest.raises(ValueError):
mmcv.flowread(np.zeros((100, 100, 1)))
def _get_flow_blob(roidb, frame_root, flow_root):
num_images = len(roidb)
scale_inds = np.random.randint(0,
high=len(cfg.TRAIN.SCALES),
size=num_images)
processed_flows = []
final_flow_scales = []
for i in range(num_images):
video_name, str_frame_index = roidb[i]['image'].split('/')
frame_index = int(str_frame_index)
flows = []
flow_scales = None
if cfg.A2D.LOAD_FLOW:
if flow_root is None or not os.path.exists(flow_root):
raise RuntimeError(
'flow root not provided or flow root not exists')
for f_idx in range(frame_index - int(cfg.A2D.SEGMENT_LENGTH / 2),
frame_index + int(cfg.A2D.SEGMENT_LENGTH / 2)):
flow_path = os.path.join(flow_root, video_name,
'%05d' % f_idx + '.flo')
assert os.path.exists(flow_path)
flow = mmcv.flowread(flow_path)
if roidb[i]['flipped']:
flow = flow[:, ::-1, :]
target_size = cfg.TRAIN.SCALES[scale_inds[i]]
flows_with_diff_scales, computed_flow_scales = blob_utils.prep_flow_for_blob(
flow,
cfg.A2D.FLOW_MAX_MAG, [target_size],
cfg.TRAIN.MAX_SIZE,
clip_mag=cfg.A2D.CLIP_FLOW_MAG)
three_channel_flow = flow_to_flow_img(
flows_with_diff_scales[0])
flows.append(three_channel_flow)
if flow_scales is not None:
assert computed_flow_scales[0] == flow_scales
else:
flow_scales = computed_flow_scales[0]
else:
prev_frame = None
for f_idx in range(
frame_index - int(cfg.A2D.SEGMENT_LENGTH / 2),
frame_index + int(cfg.A2D.SEGMENT_LENGTH / 2) + 1):
frame_fpath = os.path.join(frame_root, video_name,
'%05d' % f_idx + '.png')
assert os.path.exists(frame_fpath)
cur_frame = cv2.imread(frame_fpath)
cur_frame = cv2.cvtColor(cur_frame, cv2.COLOR_BGR2GRAY)
if prev_frame is not None:
flow = cv2.calcOpticalFlowFarneback(
prev_frame, cur_frame, None, 0.5, 3, 15, 3, 5, 1.2, 0)
if roidb[i]['flipped']:
flow = flow[:, ::-1, :]
target_size = cfg.TRAIN.SCALES[scale_inds[i]]
flows_with_diff_scales, computed_flow_scales = blob_utils.prep_flow_for_blob(
flow,
cfg.A2D.FLOW_MAX_MAG, [target_size],
cfg.TRAIN.MAX_SIZE,
clip_mag=cfg.A2D.CLIP_FLOW_MAG)
three_channel_flow = flow_to_flow_img(
flows_with_diff_scales[0])
flows.append(three_channel_flow)
if flow_scales is not None:
assert computed_flow_scales[0] == flow_scales
else:
flow_scales = computed_flow_scales[0]
prev_frame = cur_frame
processed_flows.append(flows)
final_flow_scales.append(flow_scales)
blob = blob_utils.flow_list_to_blob(processed_flows)
return blob, final_flow_scales
def validate_with_opticalflow(config,
val_loader,
val_dataset,
model,
criterion,
output_dir,
tb_log_dir,
writer_dict=None):
root = config.DATASET.ROOT
batch_time = AverageMeter()
# losses = AverageMeter()
# acc = AverageMeter()
# switch to evaluate mode
model.eval()
### load det bboxs ###
if os.path.exists(
os.path.join(
root,
'new_full_det_bboxs_' + str(config.TEST.IMAGE_THRE) + '.npy')):
print('loading new_full_det_bboxs.npy from {}...'.format(
os.path.join(
root,
'new_full_det_bboxs_' + str(config.TEST.IMAGE_THRE) + '.npy')))
full_bboxs = np.load(
os.path.join(
root, 'new_full_det_bboxs_' + str(config.TEST.IMAGE_THRE) +
'.npy')).item()
print('detection bboxes loaded')
ids = sorted(full_bboxs.keys())
else:
print('creating new_full_det_bboxs.npy...')
full_bboxs = {}
ids = []
for _, meta in val_loader:
# print(type(input))
# print(input.shape)
image_id = meta['image_id'][0].item()
if image_id not in ids:
ids.append(int(image_id))
#generate ids
ids = sorted(ids)
#fulfill the missing ids
pre_im_id = ids[0]
for im_id in ids:
if (im_id - pre_im_id) > 1 and (im_id - pre_im_id) < 60:
for i in range(im_id - pre_im_id - 1):
pre_im_id = pre_im_id + 1
if pre_im_id not in ids:
ids.append(int(pre_im_id))
logger.info(
'adding missing image_id--{}'.format(pre_im_id))
pre_im_id = im_id
ids = sorted(ids)
temp_key = {}
temp_key['ids'] = ids
# with open(os.path.join(root,'temp_id_vis.json'),'w') as f :
# json.dump(temp_key,f,indent=4)
# print('finish writing temp_id_vis.json')
for im_id in ids:
full_bboxs[im_id] = []
for _, meta in val_loader:
image_id = meta['image_id'][0].item()
center = meta['center'].numpy()
scale = meta['scale'].numpy()
score = meta['score'].numpy()
box_sc = np.array(meta['box_sc']) # [[x,y,w,h,score]]
box = (center, scale, score, box_sc)
full_bboxs[int(image_id)].append(
box) # {1003420000:[(c1,s1,score1,[x1,y1,w1,h1,score1]),
np.save(
os.path.join(
root,
'new_full_det_bboxs' + str(config.TEST.IMAGE_THRE) + '.npy'),
full_bboxs)
print('detection bboxes loaded')
with torch.no_grad():
end = time.time()
batch_time.update(time.time() - end)
image_path = []
frames = []
num_box = 0
pres, vals, c, s, sc, track_IDs = [], [], [], [], [], []
Q = deque(
maxlen=config.TEST.TRACK_FRAME_LEN) # tracked instances queue
next_track_id = FIRST_TRACK_ID
for i, im_id in enumerate(ids):
file_name = index_to_path(root, im_id)
data_numpy = cv2.imread(
file_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
frame_boxs = full_bboxs[
im_id] # get all boxes information in this frame
boxs = np.array([item[-1] for item in frame_boxs])
keep = bbox_nms(boxs, 0.5) # do the nms for each frame
if len(keep) == 0:
nmsed_boxs = frame_boxs
else:
nmsed_boxs = [frame_boxs[_keep] for _keep in keep]
print('current im_id_{}'.format(im_id))
next_id = im_id + 1
if next_id in ids:
image_id = str(im_id)
root_flow = os.path.join(root, config.TEST.FLOW_PATH)
folder_name = image_id[1:7] + '_mpii_test'
flow_name = '00' + image_id[-4:] + '.flo'
flow_path = os.path.join(root_flow, folder_name, flow_name)
flow = mmcv.flowread(flow_path) # [h,w,2]
instance = []
which_box = 0
#compute each box
for box in nmsed_boxs:
person = {}
person_flow = {}
num_box += 1
c_d = box[0]
s_d = box[1]
c_dt = box[0][0]
s_dt = box[1][0]
# print('type:{}, value:{}'.format(type(s_d),s_d))
score = box[2]
c.append(c_dt)
s.append(s_dt)
sc.append(score)
r = 0
image_path.append(file_name)
h, w = data_numpy.shape[0], data_numpy.shape[1]
trans = get_affine_transform(c_dt, s_dt, r,
config.MODEL.IMAGE_SIZE)
input = cv2.warpAffine(data_numpy,
trans,
(int(config.MODEL.IMAGE_SIZE[0]),
int(config.MODEL.IMAGE_SIZE[1])),
flags=cv2.INTER_LINEAR)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
input = transform(input)
new_input = np.zeros([
1, 3, config.MODEL.IMAGE_SIZE[1],
config.MODEL.IMAGE_SIZE[0]
])
new_input[0, :, :, :] = input[:, :, :]
input = torch.from_numpy(new_input).float()
output = model(input)
if config.TEST.FLIP_TEST:
# this part is ugly, because pytorch has not supported negative index
# input_flipped = model(input[:, :, :, ::-1])
input_flipped = np.flip(input.cpu().numpy(), 3).copy()
input_flipped = torch.from_numpy(input_flipped).cuda()
output_flipped = model(input_flipped)
output_flipped = flip_back(output_flipped.cpu().numpy(),
val_dataset.flip_pairs)
output_flipped = torch.from_numpy(
output_flipped.copy()).cuda()
if config.TEST.SHIFT_HEATMAP:
output_flipped[:, :, :, 1:] = \
output_flipped.clone()[:, :, :, 0:-1]
# output_flipped[:, :, :, 0] = 0
output = (output + output_flipped) * 0.5
batch_time.update(time.time() - end)
end = time.time()
preds, maxvals = get_final_preds(config,
output.clone().cpu().numpy(),
c_d, s_d)
#preds--(1, 17, 2)
#macvals--(1, 17, 1)
# if the value of each pred<0.4 set the joint into invisible
preds_set_invisible, maxvals_set_invisible = get_visible_joints(
preds.copy(), maxvals.copy(), config.TEST.IN_VIS_THRE)
individual = np.concatenate(
(preds_set_invisible.squeeze(),
maxvals_set_invisible.reshape(-1, 1)),
axis=1).flatten() # array with shape(num_keypoints x 3,)
person['image'] = input
person['keypoints'] = individual
# person['bbox'] = bbox[:-1]
person['score'] = score[0]
person['track_id'] = None
person['bbox'], person['area'] = get_bbox_sc_from_cs(
c_dt, s_dt, score)
instance.append(person)
pres.append(preds_set_invisible)
vals.append(maxvals_set_invisible) #[1,17,1]
#get the avg_joint_score for each box
joint_score = 0
for joint_i in range(maxvals.shape[1]):
joint_score += maxvals[0][joint_i]
avg_joint_score = joint_score / maxvals.shape[1]
#get center and scale from flow
c_flow, s_flow, box_sc_flow, is_ignore, scale_flow, flow_kps = get_cs_from_flow(
flow, preds_set_invisible, h, w, avg_joint_score,
config.TEST.FLOW_THRE) # TODO
### save debug bboxes ###
if (i % config.PRINT_FREQ
) == 0 and config.DEBUG.SAVE_ALL_BOXES:
file = data_numpy.copy()
save_all_boxes_with_joints(file, im_id, which_box, c_dt,
s_dt, c_flow[0], s_flow[0],
preds_set_invisible, score,
avg_joint_score, output_dir)
which_box += 1
if is_ignore or next_id not in ids:
continue
box_flow = (c_flow, s_flow, avg_joint_score, box_sc_flow)
full_bboxs[next_id].append(box_flow)
individual_flow = np.concatenate(
(flow_kps, maxvals.reshape(-1, 1)), axis=1).flatten()
person_flow['keypoints'] = individual_flow
person_flow['bbox'] = box_sc_flow[:-1]
person_flow['area'] = scale_flow
person_flow['track_id'] = None
person_flow['image'] = input
person_flow['score'] = score[0]
instance = instance[:-1]
instance.append(person_flow)
### Assign the Track ID for all instances in one frame ###
#when the image id is in ids but detector and flow detect nobody in the frame
if len(instance) == 0:
continue
# the last frame of a video
# go into another video
if next_id not in ids:
if config.DEBUG.SAVE_VIDEO_TRACKING:
frames = save_image_with_skeleton(data_numpy, im_id,
instance, output_dir,
frames, next_id, ids)
IDs, next_track_id = assignID(
Q,
instance,
next_track_id,
similarity_thresh=config.TEST.TRACK_SIMILARITY_THRE)
for i_person, each_person in enumerate(instance):
each_person['track_id'] = IDs[i_person]
track_IDs.extend(IDs)
Q = deque(maxlen=config.TEST.TRACK_FRAME_LEN)
next_track_id = FIRST_TRACK_ID
logger.info(
'current_im_id{}--go in to next video--next_track_id{}'.
format(im_id, next_track_id))
continue # init the Deque for the next video
IDs, next_track_id = assignID(
Q,
instance,
next_track_id,
similarity_thresh=config.TEST.TRACK_SIMILARITY_THRE)
print('IDs--{}'.format(IDs))
for i_person, each_person in enumerate(instance):
each_person['track_id'] = IDs[i_person]
#########################save image with joints and skeletons###################
if config.DEBUG.SAVE_VIDEO_TRACKING:
frames = save_image_with_skeleton(data_numpy, im_id, instance,
output_dir, frames, next_id,
ids)
track_IDs.extend(IDs)
Q.append(instance) # Q:[[{},{},{}],[{},{}],...]
#print
if i % (config.PRINT_FREQ) == 0:
msg = 'Test: [{0}/{1}]\t' \
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
.format(
i, len(ids), batch_time=batch_time,)
logger.info(msg)
logger.info('boxes number:{}\t'.format(num_box))
pres = np.array(pres)
vals = np.array(vals)
c, s, sc = np.array(c), np.array(s), np.array(sc)
np.save(os.path.join(root, 'full_pose_results.npy'), pres)
np.save(os.path.join(root, 'full_pose_scores.npy'), vals)
total_bboxes = np.zeros((num_box, 6))
total_preds = np.zeros((num_box, config.MODEL.NUM_JOINTS, 3),
dtype=np.float32) # num_box x 17 x 3
total_track_IDs = np.zeros((num_box))
for i in range(num_box):
total_preds[i:i + 1, :, 0:2] = pres[i, :, :, 0:2]
total_preds[i:i + 1, :, 2:3] = vals[i]
total_bboxes[i:i + 1, 0:2] = c[i][0:2]
total_bboxes[i:i + 1, 2:4] = s[i][0:2]
total_bboxes[i:i + 1, 4] = np.prod(s * 200, 1)[i]
total_bboxes[i:i + 1, 5] = sc[i]
total_track_IDs[i] = track_IDs[i]
name_values, perf_indicator = val_dataset.evaluate(
config, total_preds, output_dir, total_bboxes, image_path,
total_track_IDs)
return perf_indicator
warped_flow = F.grid_sample(flow, vgrid)
consistency_mask = torch.sum((warped_flow + inv_flow) ** 2, dim=1, keepdim=True) <= \
torch.sum(0.01 * (warped_flow ** 2 + inv_flow ** 2), dim=1, keepdim=True) + 0.5
warped_id_grid = F.grid_sample(id_grid, vgrid, mode='nearest')
max_id = id_grid.max()
# consistency fails
n = warped_id_grid[~consistency_mask].numel()
warped_id_grid[~consistency_mask] = torch.arange(
max_id + 1, max_id + n + 1)
max_id = max_id + n
# out of border
n = warped_id_grid[warped_id_grid == 0].numel()
warped_id_grid[warped_id_grid == 0] = torch.arange(
max_id + 1, max_id + n + 1)
max_id = max_id + n
embed()
self.id_grids = self.cat(self.id_grids, warped_id_grid)
if __name__ == '__main__':
traj = Trajectory()
flow_path = '/databack1/KITTI/kitti/tracking/training/Flow/0000/000000.flo'
inv_flow_path = '/databack1/KITTI/kitti/tracking/training/Inv_Flow/0000/000001.flo'
flow = mmcv.flowread(flow_path)
inv_flow = mmcv.flowread(inv_flow_path)
H, W, _ = flow.shape
flow = torch.tensor(flow).permute(2, 0, 1).unsqueeze(0)
inv_flow = torch.tensor(inv_flow).permute(2, 0, 1).unsqueeze(0)
frame = torch.zeros(1, 3, H, W)
traj.init(frame)
traj.update(frame, flow, inv_flow)
