Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| # @Author : xuelun | |
| import os | |
| import cv2 | |
| import torch | |
| import random | |
| import numpy as np | |
| import torch.nn.functional as F | |
| from tqdm import tqdm | |
| from os import listdir | |
| from pathlib import Path | |
| from functools import reduce | |
| from datetime import datetime | |
| from argparse import ArgumentParser | |
| from os.path import join, isdir, exists | |
| from datasets.dataset import RGBDDataset | |
| from datasets.walk import cfg | |
| from datasets.walk.utils import covision, intersected, read_images | |
| from datasets.walk.utils import fast_make_matching_robust_fitting_figure | |
| parse_mtd = lambda name: name.parent.stem.split()[1] | |
| parse_skip = lambda name: int(str(name).split(os.sep)[-1].rpartition('SP')[-1].strip().rpartition(' ')[0]) | |
| parse_resize = lambda name: str(name).split(os.sep)[-2].rpartition('[R]')[-1].rpartition('[S]')[0].strip() | |
| create_table = lambda x, y, w: dict(zip(np.round(x) + np.round(y) * w, list(range(len(x))))) | |
| class WALKDataset(RGBDDataset): | |
| def __init__(self, | |
| root_dir, # data root dit | |
| npz_root, # data info, like, overlap, image_path, depth_path | |
| seq_name, # current sequence | |
| mode, # train or val or test | |
| max_resize, # max edge after resize | |
| df, # general is 8 for ResNet w/ pre 3-layers | |
| padding, # padding image for batch training | |
| augment_fn, # augmentation function | |
| max_samples, # max sample in current sequence | |
| **kwargs): | |
| super().__init__() | |
| self.mode = mode | |
| self.root_dir = root_dir | |
| self.scene_path = join(root_dir, seq_name) | |
| pseudo_labels = kwargs.get('PSEUDO_LABELS', None) | |
| npz_paths = [join(npz_root, x) for x in pseudo_labels] | |
| npz_paths = [x for x in npz_paths if exists(x)] | |
| npz_names = [{d[:int(d.split()[-1])]: Path(path, d) for d in listdir(path) if isdir(join(path, d))} for path in npz_paths] | |
| npz_paths = [name_dict[seq_name] for name_dict in npz_names if seq_name in name_dict.keys()] | |
| self.propagating = kwargs.get('PROPAGATING', False) | |
| if self.propagating and len(npz_paths) != 24: | |
| print(f'{seq_name} has {len(npz_paths)} pseudo labels, but 24 are expected.') | |
| exit(0) | |
| self.scale = 1 / df | |
| self.scene_id = seq_name | |
| self.skips = sorted(list({parse_skip(name) for name in npz_paths})) | |
| self.resizes = sorted(list({parse_resize(name) for name in npz_paths})) | |
| self.methods = sorted(list({parse_mtd(name) for name in npz_paths}))[::-1] | |
| self.min_final_matches = kwargs.get('MIN_FINAL_MATCHES', None) | |
| self.min_filter_matches = kwargs.get('MIN_FILTER_MATCHES', None) | |
| pproot = kwargs.get('PROPAGATE_ROOT', None) | |
| ppid = ' '.join(self.methods + list(map(str, self.skips)) + self.resizes + [f'FM {self.min_filter_matches}', f'PM {self.min_final_matches}']) | |
| self.pproot = join(pproot, ppid, seq_name) | |
| if not self.propagating: | |
| assert exists(self.pproot) | |
| elif not exists(self.pproot): | |
| os.makedirs(self.pproot, exist_ok=True) | |
| image_root = kwargs.get('VIDEO_IMAGE_ROOT', None) | |
| self.image_root = join(image_root, seq_name) | |
| if not exists(self.image_root): | |
| os.makedirs(self.image_root, exist_ok=True) | |
| self.step = kwargs.get('STEP', None) | |
| self.pix_thr = kwargs.get('PIX_THR', None) | |
| self.fix_matches = kwargs.get('FIX_MATCHES', None) | |
| source_root = kwargs.get('SOURCE_ROOT', None) | |
| scap = cv2.VideoCapture(join(source_root, seq_name + '.mp4')) | |
| self.pseudo_size = [int(scap.get(3)), int(scap.get(4))] | |
| source_fps = int(scap.get(5)) | |
| video_path = join(root_dir, seq_name + '.mp4') | |
| vcap = cv2.VideoCapture(video_path) | |
| self.frame_size = [int(vcap.get(3)), int(vcap.get(4))] | |
| if self.propagating: | |
| nums = {skip: [] for skip in self.skips} | |
| idxs = {skip: [] for skip in self.skips} | |
| self.path = {skip: [] for skip in self.skips} | |
| for npz_path in npz_paths: | |
| skip = parse_skip(npz_path) | |
| assert exists(npz_path / 'nums.npy') | |
| with open(npz_path / 'nums.npy', 'rb') as f: | |
| npz = np.load(f) | |
| nums[skip].append(npz) | |
| assert exists(npz_path / 'idxs.npy') | |
| with open(npz_path / 'idxs.npy', 'rb') as f: | |
| npz = np.load(f) | |
| idxs[skip].append(npz) | |
| self.path[skip].append(npz_path) | |
| ids1 = reduce(intersected, [idxs[nums > self.min_filter_matches] for nums, idxs in zip(nums[self.skips[-1]], idxs[self.skips[-1]])]) | |
| continue1 = np.array([x in ids1[:, 0] for x in (ids1[:, 0] + self.skips[-1] * 1)]) | |
| ids2 = reduce(intersected, idxs[self.skips[-2]]) | |
| continue2 = np.array([x in ids2[:, 0] for x in ids1[:, 0]]) | |
| continue2 = continue2 & np.array([x in ids2[:, 0] for x in (ids1[:, 0] + self.skips[-2] * 1)]) | |
| ids3 = reduce(intersected, idxs[self.skips[-3]]) | |
| continue3 = np.array([x in ids3[:, 0] for x in ids1[:, 0]]) | |
| continue3 = continue3 & np.array([x in ids3[:, 0] for x in (ids1[:, 0] + self.skips[-3] * 1)]) | |
| continue3 = continue3 & np.array([x in ids3[:, 0] for x in (ids1[:, 0] + self.skips[-3] * 2)]) | |
| continue3 = continue3 & np.array([x in ids3[:, 0] for x in (ids1[:, 0] + self.skips[-3] * 3)]) | |
| continues = continue1 & continue2 & continue3 | |
| ids = ids1[continues] | |
| pair_ids = np.array(list(zip(ids[:, 0], np.clip(ids[:, 0]+self.step*self.skips[-1], a_min=ids[0, 0], a_max=ids[-1, 1])))) if self.step > 0 else ids | |
| pair_ids = pair_ids[(pair_ids[:, 1] - pair_ids[:, 0]) >= self.skips[-1]] | |
| else: | |
| pair_ids = np.array([tuple(map(int, x.split('.npy')[0].split('_'))) for x in os.listdir(self.pproot) if x.endswith('.npy')]) | |
| if (max_samples > 0) and (len(pair_ids) > max_samples): | |
| random_state = random.getstate() | |
| np_random_state = np.random.get_state() | |
| random.seed(3407) | |
| np.random.seed(3407) | |
| pair_ids = pair_ids[sorted(np.random.randint(len(pair_ids), size=max_samples))] | |
| random.setstate(random_state) | |
| np.random.set_state(np_random_state) | |
| # remove unvalid pairs from self.pproot/bad_pairs.txt | |
| pair_ids = set(map(tuple, pair_ids.tolist())) | |
| if self.propagating: | |
| assert not exists(join(self.pproot, 'bad_pairs.txt')) | |
| if exists(join(self.pproot, 'bad_pairs.txt')): | |
| with open(join(self.pproot, 'bad_pairs.txt'), 'r') as f: | |
| unvalid_pairs = set([tuple(map(int, line.split())) for line in f.readlines()]) | |
| self.unvalid_pairs_num = len(unvalid_pairs) if not self.propagating else 'N/A' | |
| pair_ids = pair_ids - unvalid_pairs | |
| self.valid_pairs_num = len(pair_ids) if not self.propagating else 'N/A' | |
| self.pair_ids = list(map(list, pair_ids)) # List[List[int, int]] | |
| # parameters for image resizing, padding and depthmap padding | |
| if mode == 'train': assert max_resize is not None | |
| self.df = df | |
| self.max_resize = max_resize | |
| self.padding = padding | |
| # for training LoFTR | |
| self.augment_fn = augment_fn if mode == 'train' else None | |
| def __len__(self): | |
| return len(self.pair_ids) | |
| def propagate(self, idx0, idx1, skips): | |
| """ | |
| Args: | |
| idx0: (int) index of the first frame | |
| idx1: (int) index of the second frame | |
| skips: (List) | |
| Returns: | |
| """ | |
| skip = skips[-1] # 40 | |
| indices = [skip * (i + 1) + idx0 for i in range((idx1 - idx0) // skip)] | |
| if (not indices) or (idx0 != indices[0]): indices = [idx0] + indices | |
| if idx1 != indices[-1]: indices = indices + [idx1] | |
| indices = list(zip(indices[:-1], indices[1:])) | |
| # [(N', 4), (N'', 4), ...] | |
| labels = [] | |
| ids = [idx0] | |
| while indices: | |
| pair = indices.pop(0) # (tuple) | |
| if pair[0] == pair[1]: break | |
| label = [] | |
| if (pair[-1] - pair[0]) == skip: | |
| tmp = self.dump(skip, pair) | |
| if len(tmp) > 0: label.append(tmp) # (ndarray) (N, 4) | |
| if skips[:-1]: | |
| _label_, id0, id1 = self.propagate(pair[0], pair[1], skips[:-1]) | |
| if (id0, id1) == pair: label.append(_label_) # (ndarray) (M, 4) | |
| if label: | |
| label = np.concatenate(label, axis=0) # (ndarray) (N+M, 4) | |
| labels.append(label) | |
| ids += [pair[1]] | |
| if len(labels) > 1: | |
| _labels_ = self.link(labels[0], labels[1]) | |
| if _labels_ is not None: | |
| labels = [_labels_] | |
| ids = [ids[0], ids[-1]] | |
| else: | |
| labels.pop(-1) | |
| ids.pop(-1) | |
| indices = [(pair[0], pair[1]-skips[0])] | |
| if len(labels) == 1 and len(ids) == 2: | |
| return labels[0], ids[0], ids[-1] | |
| else: | |
| return None, None, None | |
| def link(self, label0, label1): | |
| """ | |
| Args: | |
| label0: (ndarray) N x 4 | |
| label1: (ndarray) M x 4 | |
| Returns: (ndarray) (N', 4) | |
| """ | |
| # get keypoints in left, middle and right frame | |
| left_t0 = label0[:, :2] # (N, 2) | |
| mid_t0 = label0[:, 2:] # (N, 2) | |
| mid_t1 = label1[:, :2] # (M, 2) | |
| right_t1 = label1[:, 2:] # (M, 2) | |
| mid0_table = create_table(mid_t0[:, 0], mid_t0[:, 1], self.pseudo_size[0]) | |
| mid1_table = create_table(mid_t1[:, 0], mid_t1[:, 1], self.pseudo_size[0]) | |
| keys = {*mid0_table} & {*mid1_table} | |
| i = np.array([mid0_table[k] for k in keys]) | |
| j = np.array([mid1_table[k] for k in keys]) | |
| # remove repeat matches | |
| ij = np.unique(np.vstack((i, j)), axis=1) | |
| if ij.shape[1] < self.min_final_matches: return None | |
| # get the new pseudo labels | |
| pseudo_label = np.concatenate([left_t0[ij[0]], right_t1[ij[1]]], axis=1) # (N', 4) | |
| return pseudo_label | |
| def dump(self, skip, pair): | |
| """ | |
| Args: | |
| skip: | |
| pair: | |
| Returns: pseudo_label (N, 4) | |
| """ | |
| labels = [] | |
| for path in self.path[skip]: | |
| p = path / '{}.npy'.format(str(np.array(pair))) | |
| if exists(p): | |
| with open(p, 'rb') as f: | |
| labels.append(np.load(f)) | |
| if len(labels) > 0: labels = np.concatenate(labels, axis=0).astype(np.float32) # (N, 4) | |
| return labels | |
| def __getitem__(self, idx): | |
| idx0, idx1 = self.pair_ids[idx] | |
| pppath = join(self.pproot, '{}_{}.npy'.format(idx0, idx1)) | |
| if self.propagating and exists(pppath): | |
| return None | |
| # check propagation | |
| if not self.propagating: | |
| assert exists(pppath), f'{pppath} does not exist' | |
| if not exists(pppath): | |
| pseudo_label, idx0, idx1 = self.propagate(idx0, idx1, self.skips) | |
| if idx1 - idx0 == self.skips[-1]: | |
| pseudo_label, idx0, idx1 = self.propagate(idx0, idx1, self.skips[:-1]) | |
| if idx1 - idx0 == self.skips[-2]: | |
| pseudo_label, idx0, idx1 = self.propagate(idx0, idx1, self.skips[:-2]) | |
| if pseudo_label is None: | |
| _idx0_, _idx1_ = self.pair_ids[idx] | |
| with open(join(self.pproot, 'bad_pairs.txt'), 'a') as f: | |
| f.write('{} {}\n'.format(_idx0_, _idx1_)) | |
| return None | |
| _, mask = cv2.findFundamentalMat(pseudo_label[:, :2], pseudo_label[:, 2:], cv2.USAC_MAGSAC, ransacReprojThreshold=1.0, confidence=0.999999, maxIters=1000) | |
| mask = mask.ravel() > 0 | |
| pseudo_label = pseudo_label[mask] | |
| if len(pseudo_label) < 64 or (idx1 - idx0) == self.skips[-3]: | |
| _idx0_, _idx1_ = self.pair_ids[idx] | |
| with open(join(self.pproot, 'bad_pairs.txt'), 'a') as f: | |
| f.write('{} {}\n'.format(_idx0_, _idx1_)) | |
| return None | |
| else: | |
| with open(pppath, 'wb') as f: | |
| np.save(f, np.concatenate((np.array([[idx0, idx1, idx0, idx1]]).astype(np.float32), pseudo_label), axis=0)) | |
| else: | |
| with open(pppath, 'rb') as f: | |
| pseudo_label = np.load(f) | |
| idx0, idx1 = pseudo_label[0].astype(np.int64)[:2].tolist() | |
| pseudo_label = pseudo_label[1:] | |
| if self.propagating: | |
| return None | |
| pseudo_label *= (np.array(self.frame_size * 2) / np.array(self.pseudo_size * 2))[None] | |
| # get image | |
| img_path0 = join(self.image_root, '{}.png'.format(idx0)) | |
| color0 = cv2.imread(img_path0) | |
| img_path1 = join(self.image_root, '{}.png'.format(idx1)) | |
| color1 = cv2.imread(img_path1) | |
| width0, height0 = self.frame_size | |
| width1, height1 = self.frame_size | |
| left_upper_cornor = pseudo_label[:, :2].min(axis=0) | |
| left_low_corner = pseudo_label[:, :2].max(axis=0) | |
| left_corner = np.concatenate([left_upper_cornor, left_low_corner], axis=0) | |
| right_upper_cornor = pseudo_label[:, 2:].min(axis=0) | |
| right_low_corner = pseudo_label[:, 2:].max(axis=0) | |
| right_corner = np.concatenate([right_upper_cornor, right_low_corner], axis=0) | |
| # Prepare variables | |
| image0, color0, scale0, rands0, offset0, hlip0, vflip0, resize0, mask0 = read_images( | |
| None, self.max_resize, self.df, self.padding, | |
| np.random.choice([self.augment_fn, None], p=[0.5, 0.5]), | |
| aug_prob=1.0, is_left=True, | |
| upper_cornor=left_corner, | |
| read_size=self.frame_size, image=color0) | |
| image1, color1, scale1, rands1, offset1, hlip1, vflip1, resize1, mask1 = read_images( | |
| None, self.max_resize, self.df, self.padding, | |
| np.random.choice([self.augment_fn, None], p=[0.5, 0.5]), | |
| aug_prob=1.0, is_left=False, | |
| upper_cornor=right_corner, | |
| read_size=self.frame_size, image=color1) | |
| # warp keypoints by scale, offset and hlip | |
| pseudo_label = torch.tensor(pseudo_label, dtype=torch.float) | |
| left = (pseudo_label[:, :2] / scale0[None] - offset0[None]) | |
| left[:, 0] = resize0[1] - 1 - left[:, 0] if hlip0 else left[:, 0] | |
| left[:, 1] = resize0[0] - 1 - left[:, 1] if vflip0 else left[:, 1] | |
| right = (pseudo_label[:, 2:] / scale1[None] - offset1[None]) | |
| right[:, 0] = resize1[1] - 1 - right[:, 0] if hlip1 else right[:, 0] | |
| right[:, 1] = resize1[0] - 1 - right[:, 1] if vflip1 else right[:, 1] | |
| mask = (left[:, 0] >= 0) & (left[:, 0]*self.scale <= (resize0[1]*self.scale - 1)) & \ | |
| (left[:, 1] >= 0) & (left[:, 1]*self.scale <= (resize0[0]*self.scale - 1)) & \ | |
| (right[:, 0] >= 0) & (right[:, 0]*self.scale <= (resize1[1]*self.scale - 1)) & \ | |
| (right[:, 1] >= 0) & (right[:, 1]*self.scale <= (resize1[0]*self.scale - 1)) | |
| left, right = left[mask], right[mask] | |
| pseudo_label = torch.cat([left, right], dim=1) | |
| pseudo_label = torch.unique(pseudo_label, dim=0) | |
| fix_pseudo_label = torch.zeros(self.fix_matches, 4, dtype=pseudo_label.dtype) | |
| fix_pseudo_label[:len(pseudo_label)] = pseudo_label | |
| # read image size | |
| imsize0 = torch.tensor([height0, width0], dtype=torch.long) | |
| imsize1 = torch.tensor([height1, width1], dtype=torch.long) | |
| resize0 = torch.tensor(resize0, dtype=torch.long) | |
| resize1 = torch.tensor(resize1, dtype=torch.long) | |
| data = { | |
| # image 0 | |
| 'image0': image0, | |
| 'color0': color0, | |
| 'imsize0': imsize0, | |
| 'offset0': offset0, | |
| 'resize0': resize0, | |
| 'depth0': torch.ones((1600, 1600), dtype=torch.float), | |
| 'hflip0': hlip0, | |
| 'vflip0': vflip0, | |
| # image 1 | |
| 'image1': image1, | |
| 'color1': color1, | |
| 'imsize1': imsize1, | |
| 'offset1': offset1, | |
| 'resize1': resize1, | |
| 'depth1': torch.ones((1600, 1600), dtype=torch.float), | |
| 'hflip1': hlip1, | |
| 'vflip1': vflip1, | |
| # image transform | |
| 'pseudo_labels': fix_pseudo_label, | |
| 'gt': False, | |
| 'zs': True, | |
| # image transform | |
| 'T_0to1': torch.tensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=torch.float), | |
| 'T_1to0': torch.tensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=torch.float), | |
| 'K0': torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.float), | |
| 'K1': torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.float), | |
| # pair information | |
| 'scale0': scale0 / scale0, | |
| 'scale1': scale1 / scale1, | |
| 'rands0': rands0, | |
| 'rands1': rands1, | |
| 'dataset_name': 'WALK', | |
| 'scene_id': '{:30}'.format(self.scene_id[:min(30, len(self.scene_id)-1)]), | |
| 'pair_id': f'{idx0}-{idx1}', | |
| 'pair_names': ('{}.png'.format(idx0), | |
| '{}.png'.format(idx1)), | |
| 'covisible0': covision(pseudo_label[:, :2], resize0).item(), | |
| 'covisible1': covision(pseudo_label[:, 2:], resize1).item(), | |
| } | |
| item = super(WALKDataset, self).__getitem__(idx) | |
| item.update(data) | |
| data = item | |
| if mask0 is not None: | |
| if self.scale: | |
| # noinspection PyArgumentList | |
| [ts_mask_0, ts_mask_1] = F.interpolate(torch.stack([mask0, mask1], dim=0)[None].float(), | |
| scale_factor=self.scale, | |
| mode='nearest', | |
| recompute_scale_factor=False)[0].bool() | |
| data.update({'mask0': ts_mask_0, 'mask1': ts_mask_1}) | |
| data.update({'mask0_i': mask0, 'mask1_i': mask1}) | |
| return data | |
| if __name__ == '__main__': | |
| parser = ArgumentParser() | |
| parser.add_argument('seq_names', type=str, nargs='+') | |
| args = parser.parse_args() | |
| train_cfg = cfg.DATASET.TRAIN | |
| base_input = { | |
| 'df': 8, | |
| 'mode': 'train', | |
| 'augment_fn': None, | |
| 'max_resize': [1280, 720], | |
| 'padding': cfg.DATASET.TRAIN.PADDING, | |
| 'max_samples': cfg.DATASET.TRAIN.MAX_SAMPLES, | |
| 'min_overlap_score': cfg.DATASET.TRAIN.MIN_OVERLAP_SCORE, | |
| 'max_overlap_score': cfg.DATASET.TRAIN.MAX_OVERLAP_SCORE | |
| } | |
| cfg_input = { | |
| k: getattr(train_cfg, k) | |
| for k in [ | |
| 'DATA_ROOT', 'NPZ_ROOT', 'STEP', 'PIX_THR', 'FIX_MATCHES', 'SOURCE_ROOT', | |
| 'MAX_CANDIDATE_MATCHES', 'MIN_FINAL_MATCHES', 'MIN_FILTER_MATCHES', | |
| 'VIDEO_IMAGE_ROOT', 'PROPAGATE_ROOT', 'PSEUDO_LABELS' | |
| ] | |
| } | |
| if os.path.isfile(args.seq_names[0]): | |
| with open(args.seq_names[0], 'r') as f: | |
| seq_names = [line.strip() for line in f.readlines()] | |
| else: | |
| seq_names = args.seq_names | |
| for seq_name in seq_names: | |
| input_ = { | |
| **base_input, | |
| **cfg_input, | |
| 'root_dir': cfg_input['DATA_ROOT'], | |
| 'npz_root': cfg_input['NPZ_ROOT'], | |
| 'seq_name': seq_name | |
| } | |
| dataset = WALKDataset(**input_) | |
| random.seed(3407) | |
| np.random.seed(3407) | |
| samples = list(range(len(dataset))) | |
| num = 10 | |
| samples = random.sample(samples, num) | |
| for idx_ in tqdm(samples[:num], ncols=80, bar_format="{l_bar}{bar:3}{r_bar}", total=num, | |
| desc=f'[ {seq_name[:min(10, len(seq_name)-1)]:<10} ] [ {dataset.valid_pairs_num:<5} / {dataset.valid_pairs_num+dataset.unvalid_pairs_num:<5} ]',): | |
| data_ = dataset[idx_] | |
| if data_ is None: continue | |
| pseudo_labels_ = data_['pseudo_labels'] | |
| mask_ = pseudo_labels_.sum(dim=1) > 0 | |
| pseudo_label_ = pseudo_labels_[mask_].cpu().numpy() | |
| data_['mkpts0_f'] = pseudo_label_[:, :2] | |
| data_['mkpts1_f'] = pseudo_label_[:, 2:] | |
| data_['hw0_i'] = data_['image0'].shape[-2:] | |
| data_['hw1_i'] = data_['image1'].shape[-2:] | |
| data_['image0'] = data_['image0'][None] | |
| data_['image1'] = data_['image1'][None] | |
| data_['color0'] = data_['color0'][None] | |
| data_['color1'] = data_['color1'][None] | |
| idx0_, idx1_ = data_['pair_id'].split('-') | |
| idx0_, idx1_ = map(int, [idx0_, idx1_]) | |
| out = fast_make_matching_robust_fitting_figure(data_, transpose=True) | |
| save_dir = Path('dump/walk') / seq_name | |
| if not exists(save_dir): save_dir.mkdir(parents=True, exist_ok=True) | |
| cv2.imwrite(join(save_dir, '{:8d} [{}] {:8d} {:3d}.png'.format( | |
| idx0_, | |
| datetime.utcnow().strftime('%Y-%m-%d %H-%M-%S %f')[:-3], | |
| idx1_, | |
| idx1_ - idx0_ | |
| )), cv2.cvtColor(out, cv2.COLOR_RGB2BGR)) | |