Commit c4cf76fa authored by Clément Pinard's avatar Clément Pinard
Browse files

Add split dataset for sfmlearner

parent 3d4f804f
...@@ -414,7 +414,7 @@ All the parameters for `main_pipeline.py` are defined in the file `cli_utils.ply ...@@ -414,7 +414,7 @@ All the parameters for `main_pipeline.py` are defined in the file `cli_utils.ply
7. Ground truth creation 7. Ground truth creation
* `--eth3d_splat_radius` : Splat radius for occlusion mesh boundaries, radius of area (in meters) which will be defined as invalid because of occlusion uncertainty, see `splat_radius` option for ETH3D. Thumb rule here is that it should be around your point cloud precision. (default 0.01, i.e. 1cm) * `--eth3d_splat_radius` : Splat radius for occlusion mesh boundaries, radius of area (in meters) which will be defined as invalid because of occlusion uncertainty, see `splat_radius` option for ETH3D. Thumb rule here is that it should be around your point cloud precision. (default 0.01, i.e. 1cm)
**Tip** : As a lot of information qill go through your terminal, you can use `-vv` and `--log out.log` in order to have two terminals with different verbose levels. **Tip** : As a lot of log information will go through your terminal, you can use `-vv` and `--log out.log` in order to have two terminals with different verbose levels.
- 1st terminal : `main_pipeline` script, with `-vv` and `--log out.log` selected - 1st terminal : `main_pipeline` script, with `-vv` and `--log out.log` selected
- 2nd terminal : `tail -f out.log` - 2nd terminal : `tail -f out.log`
......
...@@ -23,6 +23,16 @@ parser.add_argument('--allow_interpolated_frames', action='store_true', ...@@ -23,6 +23,16 @@ parser.add_argument('--allow_interpolated_frames', action='store_true',
help='If set, will consider frames with interpolated odometry to be valid') help='If set, will consider frames with interpolated odometry to be valid')
def flight_path_vector(sequence, max_shift):
tvec = sequence[["pose03", "pose13", "pose23"]]
displacement = np.zeros_like(tvec)
for j in range(1, max_shift):
displacement += tvec.diff(j) / j
fpv_x = sequence["fx"] * displacement["pose03"] / displacement["pose23"] + sequence["cx"]
fpv_y = sequence["fy"] * displacement["pose13"] / displacement["pose23"] + sequence["cy"]
return fpv_x, fpv_y
def main(): def main():
args = parser.parse_args() args = parser.parse_args()
random.seed(args.seed) random.seed(args.seed)
...@@ -53,19 +63,14 @@ def main(): ...@@ -53,19 +63,14 @@ def main():
invalidity_start.append(len(valid_odometry_frames)) invalidity_start.append(len(valid_odometry_frames))
valid_sequences = [metadata.iloc[s:e].copy() for s, e in zip(validity_start, invalidity_start)] valid_sequences = [metadata.iloc[s:e].copy() for s, e in zip(validity_start, invalidity_start)]
for s in valid_sequences: for s in valid_sequences:
if len(s) <= args.min_shift: fpv_x, fpv_y = flight_path_vector(s, max_shift=3)
continue s["fpv_x"] = fpv_x
tvec = s[["pose03", "pose13", "pose23"]] s["fpv_y"] = fpv_y
displacement = np.zeros_like(tvec)
max_shift = 3
for j in range(1, max_shift):
displacement += tvec.diff(j) / j
s["fpv_x"] = s["fx"] * displacement["pose03"] / displacement["pose23"] + s["cx"]
s["fpv_y"] = s["fy"] * displacement["pose13"] / displacement["pose23"] + s["cy"]
valid_frames = s.iloc[args.min_shift:] valid_frames = s.iloc[args.min_shift:]
valid_frames = valid_frames[~valid_frames["interpolated"]] valid_frames = valid_frames[~valid_frames["interpolated"]]
valid_frames["image_path"] = [args.dataset_dir.relpathto(v) / Path(f).basename() for f in valid_frames["image_path"].values] valid_frames["image_path"] = [args.dataset_dir.relpathto(v) / Path(f).basename()
for f in valid_frames["image_path"].values]
total_valid_frames.append(valid_frames) total_valid_frames.append(valid_frames)
total_valid_frames_df = pd.concat(total_valid_frames) total_valid_frames_df = pd.concat(total_valid_frames)
......
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from path import Path
import pandas as pd
from tqdm import tqdm
import numpy as np
from scipy.spatial.transform import Rotation
import random
parser = ArgumentParser(description='Convert dataset to KITTI format, optionnally create a visualization video',
formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('--dataset_dir', metavar='DIR', type=Path, required=True,
help='folder containing the converted dataset')
parser.add_argument('--output_dir', metavar='DIR', type=Path, required=True,
help='foutput folder which will be used as a dataset for unsupervised depth training')
parser.add_argument('--scenes_list', type=Path, default=None,
help="List of folders containing videos to split, useful if some are used for testing")
parser.add_argument('--verbose', '-v', action='count', default=0)
parser.add_argument('--min_displacement', default=0, type=float,
help='Minimum displacement between two frames')
parser.add_argument('--max_rotation', default=1, type=float,
help='Maximum rotation between two frames. If above, frame is discarded and video is split')
parser.add_argument('--min_num_frames', default=5, type=int,
help='minimum number of frames in a video split. If below, split is removed')
parser.add_argument('--allow_interpolated_frames', action='store_true',
help='If set, will consider frames with interpolated odometry to be valid')
parser.add_argument('--adapt_min_disp_to_depth', action='store_true')
parser.add_argument('--train_split', type=float, default=0.9,
help='proportion of videos taken for training split, the other videos are taken for validation')
parser.add_argument('--seed', type=int, default=0)
def sample_splits(sequence, min_displacement, max_rot, min_num_frames):
def matrix(row):
flat_matrix = row[["pose00", "pose01", "pose02",
"pose10", "pose11", "pose12",
"pose20", "pose21", "pose22"]].values
rotation = Rotation.from_matrix(flat_matrix.reshape(3, 3))
return rotation
current_split = [0]
last_rot = None
last_intrinsics = None
for k in sample_frames(sequence, min_displacement):
row = sequence.iloc[k]
current_rot = matrix(row)
current_intrinsics = row[["fx", "fy", "cx", "cy"]].values.astype(float)
if last_rot is not None:
rot_diff_mag = (current_rot.inv() * last_rot).magnitude()
else:
rot_diff_mag = 0
if last_intrinsics is None:
last_intrinsics = current_intrinsics
if rot_diff_mag > max_rot or not np.allclose(current_intrinsics, last_intrinsics):
if len(current_split) > min_num_frames:
yield current_split
current_split = []
last_rot = None
last_intrinsics = None
else:
last_rot = current_rot
current_split.append(k)
if len(current_split) > min_num_frames:
yield current_split
def sample_frames(sequence, min_displacement):
tvec = sequence[["pose03", "pose13", "pose23"]].values
origin = tvec[0]
current_disp = 0
for j, pos in enumerate(tvec):
if current_disp > min_displacement:
yield j
origin = pos
current_disp = np.linalg.norm(pos - origin)
def get_min_depth(row):
depth_path = row["image_path"].splitext() + ".npy"
depth = np.load(depth_path)
return depth.min()
def sample_frames_depth(sequence, min_displacement, min_depth):
tvec = sequence[["pose03", "pose13", "pose23"]].values
origin = tvec[0]
current_disp = 0
last_min_depth = get_min_depth(sequence.iloc[0])
for i, pos in enumerate(tvec):
current_min_depth = get_min_depth(sequence.iloc[i])
if current_disp * np.min(last_min_depth, current_min_depth) / min_depth > min_displacement:
yield i
origin = pos
current_disp = np.lialg.norm(pos - origin)
def main():
args = parser.parse_args()
random.seed(args.seed)
args.output_dir.makedirs_p()
total_scenes = 0
total_frames = 0
frames_per_video = {}
dirs_per_video = {}
if args.scenes_list is not None:
with open(args.scenes_list, 'r') as f:
dataset_folders = [args.dataset_dir / path[:-1] for path in f.readlines()]
else:
dataset_folders = [f for f in args.dataset_dir.walkdirs() if f.files("*.jpg")]
for v in tqdm(dataset_folders):
frames_per_video[v] = 0
dirs_per_video[v] = []
metadata = pd.read_csv(v/"metadata.csv")
metadata["full_image_path"] = [v / Path(f).basename()
for f in metadata["image_path"].values]
valid_odometry_frames = metadata["registered"]
if not args.allow_interpolated_frames:
valid_odometry_frames = valid_odometry_frames & ~metadata["interpolated"]
# Construct valid sequences
valid_diff = valid_odometry_frames.astype(float).diff()
invalidity_start = valid_diff.index[valid_diff == -1].tolist()
validity_start = valid_diff.index[valid_diff == 1].tolist()
if valid_odometry_frames.iloc[0]:
validity_start = [0] + validity_start
if valid_odometry_frames.iloc[-1]:
invalidity_start.append(len(valid_odometry_frames))
valid_sequences = [metadata.iloc[s:e].copy() for s, e in zip(validity_start, invalidity_start)]
for i, seq in enumerate(valid_sequences):
splits = sample_splits(seq, args.min_displacement, args.max_rotation, args.min_num_frames)
for j, s in enumerate(splits):
total_scenes += 1
total_frames += len(s)
frames_per_video[v] += len(s)
final_seq = seq.iloc[s]
relative_dir = args.dataset_dir.relpathto(v)
folder_tree = relative_dir.splitall()[1:]
folder_name = '_'.join(folder_tree) + '_{:02d}_{:02d}'.format(i, j)
output_dir = args.output_dir / folder_name
dirs_per_video[v].append(folder_name)
output_dir.makedirs_p()
for _, row in final_seq.iterrows():
image_name = Path(row["image_path"]).basename()
src_img_path = v / image_name
tgt_img_path = output_dir / image_name
src_depth_path = src_img_path.stripext() + '.npy'
tgt_depth_path = tgt_img_path.stripext() + '.npy'
src_img_path.copy(tgt_img_path)
src_depth_path.copy(tgt_depth_path)
poses = final_seq[["pose00", "pose01", "pose02", "pose03",
"pose10", "pose11", "pose12", "pose13",
"pose20", "pose21", "pose22", "pose23"]]
np.savetxt(output_dir / "poses.txt", poses)
final_seq.to_csv(output_dir / "metadata.csv")
intrinsics = final_seq[["fx", "fy", "cx", "cy"]].values[0]
cam = np.array([[intrinsics[0], 0, intrinsics[2]],
[0, intrinsics[1], intrinsics[3]],
[0, 0, 1]])
np.savetxt(output_dir / "cam.txt", cam)
print("constructed a training set of {} frames in {} scenes".format(total_frames, total_scenes))
random.shuffle(dataset_folders)
train_frames = total_frames * args.train_split
cumulative_num_frames = np.cumsum(np.array([frames_per_video[v] for v in dataset_folders]))
used_for_train = (cumulative_num_frames - train_frames) < 0
train_folders = [dirs_per_video[dataset_folders[i]] for i in np.where(used_for_train)[0]]
val_folders = [dirs_per_video[dataset_folders[i]] for i in np.where(~used_for_train)[0]]
with open(args.output_dir / "train.txt", 'w') as f:
f.writelines([scene + '\n' for tf in train_folders for scene in tf])
with open(args.output_dir / "val.txt", 'w') as f:
f.writelines([scene + '\n' for vf in val_folders for scene in vf])
if __name__ == '__main__':
main()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment