Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,295 Bytes
5a0778e |
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 |
import os
import numpy as np
import os.path as osp
import json
import glob
import cv2
import shutil
from PIL import Image
from natsort import natsorted
def even_or_odd(num):
if num % 2 == 0:
return num
else:
return num - 1
def gen_json(root_path, dataset, start_id, end_id, step, save_path=None):
rgb_name = "rgb"
if dataset == "kitti":
factor = 256.0
elif dataset == "nyuv2":
factor = 6000.0
elif dataset == "bonn":
factor = 5000.0
elif dataset == 'sintel':
factor = 65535 / 650
rgb_name = "clean"
elif dataset == 'scannet':
factor = 1000.0
rgb_name = "color"
else:
raise NotImplementedError
data = {}
data[dataset] = []
pieces = glob.glob(osp.join(root_path, "*"))
count = 0
for piece in pieces:
if not osp.isdir(piece):
continue
name = piece.split('/')[-1]
name_dict = {name:[]}
images = glob.glob(osp.join(piece, rgb_name, "*.png")) + glob.glob(osp.join(piece, rgb_name, "*.jpg"))
images = natsorted(images)
depths = glob.glob(osp.join(piece, "depth/*.png"))
depths = natsorted(depths)
images = images[start_id:end_id:step]
depths = depths[start_id:end_id:step]
for i in range(len(images)):
image = images[i]
xx = image[len(root_path)+1:]
depth = depths[i][len(root_path)+1:]
tmp = {}
tmp["image"] = xx
tmp["gt_depth"] = depth
tmp["factor"] = factor
name_dict[name].append(tmp)
data[dataset].append(name_dict)
with open(save_path, "w") as f:
json.dump(data, f, indent= 4)
def gen_json_scannet_tae(root_path, start_id, end_id, step, save_path=None):
data = {}
data["scannet"] = []
pieces = glob.glob(osp.join(root_path, "*"))
color = 'color_origin'
for piece in pieces:
if not osp.isdir(piece):
continue
name = piece.split('/')[-1]
name_dict = {name:[]}
images = glob.glob(osp.join(piece,color, "*.jpg"))
images = natsorted(images)
depths = glob.glob(osp.join(piece, "depth/*.png"))
depths = natsorted(depths)
images = images[start_id:end_id:step]
depths = depths[start_id:end_id:step]
print(f"sequence frame number: {piece}")
count = 0
for i in range(len(images)):
image = images[i]
xx = image[len(root_path)+1:]
depth = depths[i][len(root_path)+1:]
base_path = osp.dirname(image)
base_path = base_path.replace(color, 'intrinsic')
K = np.loadtxt(base_path + '/intrinsic_depth.txt')
pose_path = image.replace(color, 'pose').replace('.jpg', '.txt')
pose = np.loadtxt(pose_path)
tmp = {}
tmp["image"] = xx
tmp["gt_depth"] = depth
tmp["factor"] = 1000.0
tmp["K"] = K.tolist()
tmp["pose"] = pose.tolist()
name_dict[name].append(tmp)
data["scannet"].append(name_dict)
with open(save_path, "w") as f:
json.dump(data, f, indent= 4)
def get_sorted_files(root_path, suffix):
all_img_names = os.listdir(root_path)
all_img_names = [x for x in all_img_names if x.endswith(suffix)]
print(f"sequence frame number: {len(all_img_names)}")
all_img_names.sort()
all_img_names = sorted(all_img_names, key=lambda x: int(x.split(".")[0][-4:]))
return all_img_names
def copy_crop_files(im_path, depth_path, out_img_path, out_depth_path, dataset):
img = np.array(Image.open(im_path))
if dataset == "kitti" or dataset == "bonn":
height, width = img.shape[:2]
height = even_or_odd(height)
width = even_or_odd(width)
img = img[:height, :width]
elif dataset == "nyuv2":
img = img[45:471, 41:601, :]
elif dataset == "scannet":
img = img[8:-8, 11:-11, :]
os.makedirs(osp.dirname(out_img_path), exist_ok=True)
os.makedirs(osp.dirname(out_depth_path), exist_ok=True)
cv2.imwrite(
out_img_path,
img,
)
shutil.copyfile(depth_path, out_depth_path)
|