Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add script to visualize BOP files #281

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add script to visualize BOP files #281

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
371 changes: 371 additions & 0 deletions scripts/metrics/render_bop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,371 @@

import nvisii as visii
import argparse
import numpy as np

import simplejson as json
import pyrender
import cv2
import numpy as np
import os
import glob
import pyrr
import copy

parser = argparse.ArgumentParser()

parser.add_argument(
'--spp',
default=100,
type=int,
help = "number of sample per pixel, higher the more costly"
)
parser.add_argument(
'--opencv',
action='store_true',
default=False,
help = "Use the opencv coordinate frame, e.g., dope inference,\
otherwise it uses the opengl coordinate frame."
)

parser.add_argument(
'--contour',
action='store_true',
default=False,
help = "Only draw the contour instead of the 3d model overlay"
)

parser.add_argument(
'--overlay',
action='store_true',
default=False,
help = "add the overlay"
)

parser.add_argument(
'--gray',
action='store_true',
default=False,
help = "draw the 3d model in gray"
)

parser.add_argument(
'--path_json',
required=True,
help = "path to the json files you want loaded,\
it assumes that there is a png accompanying."
)

parser.add_argument(
'--objs_folder',
default='content/',
help = "object to load folder, should follow YCB structure"
)

parser.add_argument(
'--scale',
default=1,
type=float,
help='Specify the scale of the target object(s). If the obj mesh is in '
'meters -> scale=1; if it is in cm -> scale=0.01.'
)

parser.add_argument(
'--out',
default='overlay.png',
help = "output filename"
)

opt = parser.parse_args()

# # # # # # # # # # # # # # # # # # # # # # # # #

def create_obj(
name = 'name',
path_obj = "",
path_tex = None,
scale = 1,
rot_base = None, #visii quat
pos_base = (-10,-10,-10), # visii vec3
):


# This is for YCB like dataset
if path_obj in create_obj.meshes:
obj_mesh = create_obj.meshes[path_obj]
else:
obj_mesh = visii.mesh.create_from_obj(name, path_obj)
create_obj.meshes[path_obj] = obj_mesh

try:
visii.transform.remove(name)
visii.material.remove(name)
visii.entity.remove(name)
except:
pass

obj_entity = visii.entity.create(
name = name,
# mesh = visii.mesh.create_sphere("mesh1", 1, 128, 128),
mesh = obj_mesh,
transform = visii.transform.create(name),
material = visii.material.create(name)
)

# should randomize
obj_entity.get_material().set_metallic(0) # should 0 or 1
obj_entity.get_material().set_transmission(0) # should 0 or 1
obj_entity.get_material().set_roughness(1) # default is 1

if not path_tex is None:

if path_tex in create_obj.textures:
obj_texture = create_obj.textures[path_tex]
else:
obj_texture = visii.texture.create_from_image(name,path_tex)
create_obj.textures[path_tex] = obj_texture


obj_entity.get_material().set_base_color_texture(obj_texture)

obj_entity.get_transform().set_scale(visii.vec3(scale))

if not rot_base is None:
obj_entity.get_transform().set_rotation(rot_base)
if not pos_base is None:
obj_entity.get_transform().set_position(pos_base)
print(f' created: {obj_entity.get_name()}')
return obj_entity

create_obj.meshes = {}
create_obj.textures = {}

# # # # # # # # # # # # # # # # # # # # # # # # #

visii.initialize_headless()
visii.enable_denoiser()

camera = visii.entity.create(
name = "camera",
transform = visii.transform.create("camera"),
camera = visii.camera.create(
name = "camera",
)
)

camera.get_transform().look_at(
visii.vec3(0,0,-1), # look at (world coordinate)
visii.vec3(0,1,0), # up vector
visii.vec3(0,0,0), # camera_origin
)

visii.set_camera_entity(camera)

visii.set_dome_light_intensity(1)

try:
visii.set_dome_light_color(visii.vec3(1, 1, 1), 0)
except TypeError:
# Support for alpha transparent backgrounds was added in nvisii ef1880aa,
# but as of 2022-11-03, the latest released version (1.1) does not include
# that change yet.
print("WARNING! Your version of NVISII does not support alpha transparent backgrounds yet; --contour will not work properly.")
visii.set_dome_light_color(visii.vec3(1, 1, 1))

# # # # # # # # # # # # # # # # # # # # # # # # #

# LOAD THE SCENE

objects_added = []
dope_trans = []
gt = None

with open(opt.path_json) as f:
data_json = json.load(f)

first_run = True

scene_gt = copy.deepcopy(data_json)

for frame, scene_objs in scene_gt.items():
camera_data = {}

with open(opt.path_json.replace("gt","camera")) as f:
camera_data_json = json.load(f)

camera_data['intrinsics'] = {}
camera_data['intrinsics']['fx'] = camera_data_json[frame]["cam_K"][0]
camera_data['intrinsics']['fy'] = camera_data_json[frame]["cam_K"][4]
camera_data['intrinsics']['cx'] = camera_data_json[frame]["cam_K"][2]
camera_data['intrinsics']['cy'] = camera_data_json[frame]["cam_K"][5]

# load an image for resolution
zfill_num = 1
while not os.path.exists(os.path.join(os.path.dirname(opt.path_json), "rgb", f"{frame.zfill(zfill_num)}.png")) and zfill_num < 10:
zfill_num += 1

img_path = os.path.join(os.path.dirname(opt.path_json), "rgb", f"{frame.zfill(zfill_num)}.png")
print(f"Reading image from {img_path}")
opt.im_path = img_path

img = cv2.imread(img_path)
camera_data['height'] = img.shape[0]
camera_data['width'] = img.shape[1]

# set the camera
intrinsics = camera_data['intrinsics']
im_height = camera_data['height']
im_width = camera_data['width']

cam = pyrender.IntrinsicsCamera(intrinsics['fx'],intrinsics['fy'],intrinsics['cx'],intrinsics['cy'])

proj_matrix = cam.get_projection_matrix(im_width, im_height)

proj_matrix = visii.mat4(
proj_matrix.flatten()[0],
proj_matrix.flatten()[1],
proj_matrix.flatten()[2],
proj_matrix.flatten()[3],
proj_matrix.flatten()[4],
proj_matrix.flatten()[5],
proj_matrix.flatten()[6],
proj_matrix.flatten()[7],
proj_matrix.flatten()[8],
proj_matrix.flatten()[9],
proj_matrix.flatten()[10],
proj_matrix.flatten()[11],
proj_matrix.flatten()[12],
proj_matrix.flatten()[13],
proj_matrix.flatten()[14],
proj_matrix.flatten()[15],
)
proj_matrix = visii.transpose(proj_matrix)

camera.get_camera().set_projection(proj_matrix)

all_objs = []

# get the objects to load.
for obj in scene_objs:
print(f"obj: {obj}")
to_add = {}
to_add['class'] = str(obj['obj_id'])
to_add['location'] = obj['cam_t_m2c']

# figuring out the rotation now
rot = obj['cam_R_m2c']
m = pyrr.Matrix33(
[
[rot[0],rot[1],rot[2]],
[rot[3],rot[4],rot[5]],
[rot[6],rot[7],rot[8]],
])
quat = m.quaternion
to_add['quaternion_xyzw'] = [quat.x,quat.y,quat.z,quat.w]
all_objs.append(to_add)


for i_obj, obj in enumerate(all_objs):
name = obj['class']

# hack because notation is old
if name == '003':
name = '003_cracker_box_16k'

print('loading',name)
entity_visii = create_obj(
name = obj['class'] + "_" + str(i_obj),
path_obj = f"{opt.objs_folder}/obj_{str(name).zfill(6)}.ply",
path_tex = None, # f"{opt.objs_folder}/obj_{str(name).zfill(6)}.png",
scale = .001,
rot_base = None
)

first_run = False

pos = obj['location']
rot = obj['quaternion_xyzw']

entity_visii.get_transform().set_rotation(
visii.quat(
rot[3],
rot[0],
rot[1],
rot[2],
)
)

entity_visii.get_transform().set_position(
visii.vec3(
pos[0]/1000,
pos[1]/1000,
pos[2]/1000,
)
)

entity_visii.get_transform().rotate_around(visii.vec3(0,0,0),visii.angleAxis(visii.pi(), visii.vec3(1,0,0)))

# # # # # # # # # # # # # # # # # # # # # # # # #
output_path = os.path.join(os.path.dirname(opt.path_json), "overlay", f"{frame.zfill(6)}.png")

os.makedirs(os.path.dirname(output_path), exist_ok=True)

visii.render_to_file(
width=im_width,
height=im_height,
samples_per_pixel=400,
file_path=output_path
)
# raise()

# # # # # # # # # # # # # # # # # # # # # # # # #

# # create overlay

im_path = opt.im_path
if os.path.exists(opt.path_json.replace("json",'png')):
im_path = opt.path_json.replace("json",'png')

elif os.path.exists(opt.path_json.replace("json",'jpg')):
im_path = opt.path_json.replace("json",'jpg')

im = cv2.imread(im_path)

if im is not None:
im_pred = cv2.imread(output_path, cv2.IMREAD_UNCHANGED)

alpha = im_pred[:,:,-1]/255.0 * 0.75
alpha = np.stack((alpha,)*3, axis=-1)
im_pred = im_pred.astype(float)

im = im.astype(float)

foreground = cv2.multiply(alpha,im_pred[:,:,:3])


if opt.gray:
im_gray = cv2.imread(im_path)
im_gray = cv2.cvtColor(im_gray, cv2.COLOR_BGR2GRAY)
im_gray = im_gray.astype(float)
im_gray = np.stack((im_gray,)*3, axis=-1)
background = cv2.multiply(1-alpha,im_gray)
elif opt.overlay:
background = cv2.multiply(1-alpha,im[:,:,:3])
else:
background = im[:,:,:3]
foreground = background

outrgb = cv2.add(foreground,background)

if opt.contour:
gray = cv2.cvtColor((alpha*255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
cnts = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]

for c in cnts:
cv2.drawContours(outrgb, [c], -1, (36, 255, 12), thickness=3,lineType=cv2.LINE_AA)
print(f'outputing {output_path}')
cv2.imwrite(output_path, outrgb)

# let's clean up the GPU
visii.deinitialize()