Commit 93a9955d authored by nviolante's avatar nviolante
Browse files

fix annotate dataset

parent 89bf1aae
......@@ -2,6 +2,7 @@ import click
import json
from PIL import Image
import os
import numpy as np
import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.ops import box_area
......@@ -41,7 +42,8 @@ def is_image_ext(fname):
)
@click.option("--max-images", type=int, default=None)
@click.option("--batch-size", type=int, default=32)
def main(cfg, data, dest, stylegan_format, max_images, batch_size):
@click.option("--save-record", type=bool, default=False)
def main(cfg, data, dest, stylegan_format, max_images, batch_size, save_record):
intrinsics = [
4.2647,
0.0,
......@@ -60,6 +62,9 @@ def main(cfg, data, dest, stylegan_format, max_images, batch_size):
torch.backends.cudnn.deterministic = cfg["cudnn"]["deterministic"]
torch.backends.cudnn.enabled = cfg["cudnn"]["enabled"]
if save_record:
os.makedirs("./records", exist_ok=True)
# Initialize Ego-Net and load the pre-trained checkpoint
print("Loading models..")
model = maskrcnn_resnet50_fpn(pretrained=True).eval().cuda()
......@@ -107,6 +112,12 @@ def main(cfg, data, dest, stylegan_format, max_images, batch_size):
else:
name = image_name
dataset_json["labels"].append([name, cam2world + intrinsics])
if save_record:
record['cam2world'] = cam2world
record['rotation'] = rotation
np.save(f"./camera_calib/results/{image_name.split('.')[0]}.npy", record)
total += 1
print(f"Processing {total} images")
......
......@@ -132,12 +132,12 @@ def create_canonical_box(l=2, h=1, w=1, prediction=None):
np.sum(lines[8:]) / 4,
)
x_corners = [l, l, l, l, 0, 0, 0, 0]
x_corners = [w, w, 0, 0, w, w, 0, 0]
y_corners = [0, h, 0, h, 0, h, 0, h]
z_corners = [w, w, 0, 0, w, w, 0, 0]
x_corners -= np.float32(l) / 2.0
z_corners = [l, l, l, l, 0, 0, 0, 0]
x_corners -= np.float32(w) / 2.0
y_corners -= np.float32(h) / 2.0
z_corners -= np.float32(w) / 2.0
z_corners -= np.float32(l) / 2.0
corners_3d = np.array([x_corners, y_corners, z_corners])
return corners_3d.T
......@@ -188,18 +188,6 @@ def add_coords(name: str, coord_system: CoordSystem):
)
def create_egonet2eg3d():
egonet2eg3d = np.zeros((4, 4), dtype=np.float32)
roty = np.zeros((4, 4), dtype=np.float32)
rotx = np.zeros((4, 4), dtype=np.float32)
roty[:3, :3] = Rotation.from_rotvec([0, 90, 0], degrees=True).as_matrix()
roty[3, 3] = 1.0
rotx[:3, :3] = Rotation.from_rotvec([-90, 0, 0], degrees=True).as_matrix()
rotx[3, 3] = 1.0
egonet2eg3d = rotx @ roty
return egonet2eg3d
def transform_points(object_sys, points, transformation):
to_world = np.eye(4)
to_world[:3, -1] = -object_sys.origin
......@@ -225,10 +213,8 @@ def generate_cam2world(rotation):
rotation = np.eye(4)
rotation[:3, :3] = cv.Rodrigues(rot_vec)[0]
cam = transform_coord_system(eg3d_object_sys, rotation)
# convert eg3d object sys to egd camera sys
cam2object = np.eye(4)
eg3d_camera = transform_coord_system(eg3d_object_sys, cam2object)
eg3d_camera = eg3d_object_sys
# apply rotation in eg3d camera sys
t = np.eye(4)
......@@ -258,15 +244,14 @@ if __name__ == "__main__":
# Load cars annoations
for num in nums:
record = np.load(
f"./camera_calib/results/img0000{num}_at8.npy", allow_pickle=True
f"./camera_calib/results/img0000{num}.npy", allow_pickle=True
).item()
rotation = record["rotation"]
cam2world = generate_cam2world(rotation)
cam2world = np.array(record['cam2world']).reshape(4, 4)
cam = transform_coord_system(eg3d_object_sys, cam2world)
cams.append(cam)
# Load reference cat annotations
canonical_box = ObjectBox(eg3d_object_sys, create_canonical_box(), faces, "box")
canonical_box = ObjectBox(eg3d_object_sys, create_canonical_box(3, 1.5, 1.5), faces, "box")
cam2world = np.array(
[
0.9931507213317242,
......
Supports Markdown
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