def run_custom():
from tools import handle_custom_dataset
if args.test:
#data_root = 'data/syn_colibri_v1_test'
data_root = 'data/real_colibri_v1_val'
#data_root = 'data/custom-test'
else:
data_root = 'data/real_colibri_v1_train'
#data_root = 'data/custom-train'
handle_custom_dataset.sample_fps_points(data_root)
handle_custom_dataset.custom_to_coco(data_root)
def run_custom():
from tools import handle_custom_dataset
data_root = 'data/custom'
handle_custom_dataset.sample_fps_points(data_root)
handle_custom_dataset.custom_to_coco(data_root)
def main(args):
if not os.path.isfile(args.dataset):
print(f"Invalid dataset file: {args.dataset}")
exit(1)
if not (os.path.isfile(args.model)
and os.path.splitext(args.model)[1] == ".ply"):
print(f"Invalid model path: {args.model}")
exit(1)
if not os.path.isdir(args.output):
print(f"Invalid output path: {args.output}")
exit(1)
# Create output directory structure
output_root = os.path.join(args.output, args.name)
print(f"Storing formatted dataset in {output_root}")
os.mkdir(output_root)
os.mkdir(os.path.join(output_root, "rgb"))
os.mkdir(os.path.join(output_root, "mask"))
os.mkdir(os.path.join(output_root, "pose"))
# Get source dataset root directory
dataset_root = os.path.dirname(args.dataset)
# Extract object diameter
model = trimesh.load(args.model)
diameter = max_distance(model)
with open(os.path.join(output_root, "diameter.txt"), "w") as f:
f.write(str(diameter))
# Copy model file
copyfile(args.model, os.path.join(output_root, "model.ply"))
idx = 0
calib = None
sample_names = []
with open(args.dataset, "r") as f:
sample_names = f.read().splitlines()
for sample_name in tqdm(sample_names):
# Check meta file
metaFile = os.path.join(dataset_root, f"meta/{sample_name}.pkl")
if not os.path.isfile(metaFile):
print(
f"Could not find meta file for: {sample_name}\nLooked for {metaFile}"
)
continue
meta = np.load(metaFile, allow_pickle=True)
# Check RGB image
rgbFile = os.path.join(dataset_root, f"rgb/{sample_name}.jpg")
if not os.path.isfile(rgbFile):
print(
f"Could not find rgb file for: {sample_name}\nLooked for {rgbFile}"
)
continue
# Read object mask
maskFile = os.path.join(dataset_root, f"segm/{sample_name}.png")
if not os.path.isfile(maskFile):
print(
f"Could not find mask file for: {sample_name}\nLooked for {maskFile}"
)
else:
src = cv2.imread(maskFile, cv2.IMREAD_UNCHANGED)
mask = src[:, :, 2]
if np.all(mask == 0):
print(f"Found empty mask file: {maskFile}")
continue
# Write (or check) camera parameters
if calib is None:
calib = meta['cam_calib']
np.savetxt(os.path.join(output_root, "camera.txt"), calib)
elif np.any(calib != meta['cam_calib']):
print(
"Camera calibration matrix changed within dataset, which is currently not supported! Skipping file."
)
print(f"Calib: {calib}")
print(f"Calib2: {meta['cam_calib']}")
continue
# Copy RGB images
copyfile(rgbFile, os.path.join(output_root, f"rgb/{idx}.jpg"))
# Save object mask
cv2.imwrite(os.path.join(output_root, f"mask/{idx}.png"), mask)
# Save pose information
pose = meta['cam_extr'] @ meta['affine_transform']
np.save(os.path.join(output_root, f"pose/pose{idx}.npy"), pose)
idx = idx + 1
print(f"Copied {idx} samples to new dataset.")
handle_custom_dataset.sample_fps_points(output_root)
handle_custom_dataset.custom_to_coco(output_root)
print("Done!")