def load_image(data_mean, data_std, device, image_name):
"""
Helper method to load an image into a torch tensor. Includes preprocessing.
"""
im = Image.open(image_name)
x = Normalize(mean=data_mean,
std=data_std)(ToTensor()(CenterCrop(224)(Resize(256)(im))))
x = x.unsqueeze(0).to(device)
return x
def __call__(self, vis):
global state
# Do nothing
if self.key == 0:
pass
# Dump current output window
elif self.key == ord('X'):
args.dump_dir.mkdir(exist_ok=True)
pascal_az = (int(state['angle_z']) + 90) % 360
pascal_el = 90 - int(state['angle_y'])
rad = int(state['radius'])
d_id = state['dump_id']
id_str = f'{d_id:03d}_el_{pascal_el:03d}_az_{pascal_az:03d}_rad_{rad:03d}'
dump_image_path = str(args.dump_dir / f'{id_str}.png')
cv2.imwrite(dump_image_path, state['dump_image'])
print(f'Saved {dump_image_path}.')
state['dump_id'] += 1
# Override open3D reset
elif self.key == ord('R'):
pass
# Rotation around Y axis
elif self.key == ord('F'):
state['angle_y'] += 5
elif self.key == ord('D'):
state['angle_y'] -= 5
# Rotation around Z axis
elif self.key == ord('S'):
state['angle_z'] += 5
elif self.key == ord('A'):
state['angle_z'] -= 5
# Distance from origin (+)
elif self.key == ord('H'):
state['radius'] += 0.05
# Distance from origin (-)
elif self.key == ord('G'):
state['radius'] -= 0.05
# Next dataset example
elif self.key == ord(' '):
state['texture_src'] = state['dataset'][state['dataset_index']]
state['dataset_index'] += 1
# Next CAD model
elif self.key == ord('N'):
state['cad_idx'] += 1
if state['cad_idx'] == 10:
state['cad_idx'] = 0
# Update model and 3D keypoints
cad_idx = state['cad_idx']
model_path = args.CAD_root / f'pascal_{state["pascal_class"]}_cad_{cad_idx:03d}.ply'
# Load 3D keypoints for current model
yaml_file = model_path.parent / (model_path.stem + '.yaml')
state['kpoints_3d'] = load_yaml_file(yaml_file)['kpoints_3d']
mesh = o3d.read_triangle_mesh(str(model_path))
# Compute normal colors
mesh.compute_vertex_normals()
state['normal_vertex_colors'] = (np.asarray(mesh.vertex_normals) +
1) / 2.
if 'mesh' not in state['geometries']:
state['geometries']['mesh'] = mesh
else:
state['geometries']['mesh'].vertices = mesh.vertices
state['geometries']['mesh'].vertex_colors = mesh.vertex_colors
state['geometries'][
'mesh'].vertex_normals = mesh.vertex_normals
state['geometries']['mesh'].triangles = mesh.triangles
else:
raise NotImplementedError()
# Set normal colors to the mesh
state['geometries']['mesh'].vertex_colors = o3d.Vector3dVector(
state['normal_vertex_colors'])
# Move Camera
angle_y = np.clip(state['angle_y'], -90, 90)
radius = np.clip(state['radius'], 0, state['radius'])
pascal_az = (state['angle_z'] + 90) % 360
pascal_el = 90 - angle_y
intrinsic = intrinsic_matrix(state['focal'],
cx=img_w / 2,
cy=img_h / 2)
if args.verbose:
print(f'Azimuth:{pascal_az} Elevation:{angle_y} Radius:{radius}')
extrinsic = pascal_vpoint_to_extrinsics(az_deg=pascal_az,
el_deg=pascal_el,
radius=radius)
if not vis.get_render_option() or not vis.get_view_control():
vis.update_geometry() # we don't have anything, return
return
align_view(vis, focal=state['focal'], extrinsic=extrinsic)
vis.get_render_option().mesh_color_option = o3d.MeshColorOption.Color
vis.get_render_option().light_on = False
vis.get_render_option().background_color = (0, 0, 0)
# Capture normal 2.5D sketch
src_normal = np.asarray(
vis.capture_screen_float_buffer(do_render=True))
src_normal = (src_normal * 255).astype(np.uint8)
object_mask = np.all(src_normal == 0, axis=-1)
if args.LAB:
src_normal = cv2.cvtColor(src_normal, cv2.COLOR_RGB2LAB)
else:
raise ValueError('Released model was trained in LAB space.')
# Project model kpoints in 2D
kpoints_2d_step_dict = {}
for k_name, k_val in state['kpoints_3d'].items():
point_3d = np.asarray([k_val])
kpoints_2d_step = project_points(point_3d, intrinsic, extrinsic)
kpoints_2d_step /= (img_w, img_h)
kpoints_2d_step = np.clip(kpoints_2d_step, -1, 1)
kpoints_2d_step_dict[k_name] = kpoints_2d_step.squeeze(0)
meta = {
'kpoints_2d': kpoints_2d_step_dict,
'vpoint': [pascal_az, pascal_el],
'cad_idx': state['cad_idx']
}
dst_pl_info = get_planes(np.zeros((img_h, img_w, 3)),
meta=meta,
pascal_class=args.pascal_class,
vis_oracle=state['vis_oracle'])
_, dst_kpoints_planes, dst_visibilities = dst_pl_info
texture_src = state['texture_src']
src_planes = np.asarray(
[to_image(i, from_LAB=args.LAB) for i in texture_src['planes']])
src_kpoints_planes = texture_src['src_kpoints_planes']
src_visibilities = texture_src['src_vs']
planes_warped, planes_unwarped = warp_unwarp_planes(
src_planes=src_planes,
src_planes_kpoints=src_kpoints_planes,
dst_planes_kpoints=dst_kpoints_planes,
src_visibilities=src_visibilities,
dst_visibilities=dst_visibilities,
pascal_class=args.pascal_class)
planes_warped = TextureDatasetWithNormal.planes_to_torch(
planes_warped, to_LAB=args.LAB)
planes_warped = planes_warped.reshape(
1, planes_warped.shape[0] * planes_warped.shape[1],
planes_warped.shape[2], planes_warped.shape[3])
src_sketch_input = Normalize(mean=[0.5] * 3, std=[0.5] * 3)(ToTensor()(
(Image.fromarray(src_normal))))
src_central = texture_src['src_central']
gen_in_src = torch.cat([
src_sketch_input.unsqueeze(0),
src_central.unsqueeze(0), planes_warped
],
dim=1).to(args.device)
net_image = to_image(state['net'](gen_in_src)[0], from_LAB=args.LAB)
# Use the normal image to mask artifacts
net_image[object_mask] = 255
out_image = np.concatenate([
to_image(src_sketch_input, from_LAB=args.LAB),
to_image(src_central, from_LAB=args.LAB), net_image,
to_image(texture_src['src_image'], from_LAB=args.LAB)
],
axis=1)
state['dump_image'] = out_image
cv2.imshow('Output', out_image)
cv2.waitKey(20)
vis.update_geometry()
from PIL import Image
model = EESPNet(args)
model.eval()
weigth_dict = torch.load(args.weights)
model.load_state_dict(weigth_dict)
rgb_img = Image.open('../train_140_im/MP_0001.tiff.jpg').convert('RGB')
from torchvision.transforms import CenterCrop, ToTensor, Normalize, Resize
from data_loader.melanoma_classification import MEAN, STD
crop_image = Resize(size=args.im_size)(rgb_img)
crop_image.save('../temp/rgb.jpg')
crop_image = ToTensor()(crop_image)
crop_image_norm = Normalize(mean=MEAN, std=STD)(crop_image)
crop_image_norm = crop_image_norm.unsqueeze(0)
_, outputs = model(crop_image_norm)
import matplotlib.pyplot as plt
for i in range(len(outputs)):
pred = outputs[i]
pred = F.softmax(pred, dim=1).squeeze(0)
pred = pred.max(0)[0]
h, w = pred.size()
pred = pred.cpu().detach().numpy()
plt.imshow(pred, cmap='hot', interpolation='nearest')
plt.savefig('../temp/pred_{}_{}.png'.format(h, w))
def __call__(self, vis):
global state
global kpoints_3d
global normal_vertex_colors
global part_vertex_colors
global texture_src
# Do nothing
if self.key == 0:
pass
# Dump current output window
elif self.key == ord('X'):
args.dump_dir.mkdir(exist_ok=True)
el = int(state['angle_y'])
az = int(state['angle_z'])
rad = int(state['radius'])
d_id = state['dump_id']
id_str = f'{d_id:03d}_el_{el:03d}_az_{az:03d}_rad_{rad:03d}'
cv2.imwrite(str(args.dump_dir / f'{id_str}.png'),
state['dump_image'])
state['dump_id'] += 1
# Override open3D reset
elif self.key == ord('R'):
pass
# Rotation around Y axis
elif self.key == ord('F'):
state['angle_y'] += 5
elif self.key == ord('D'):
state['angle_y'] -= 5
# Rotation around Z axis
elif self.key == ord('S'):
state['angle_z'] += 5
elif self.key == ord('A'):
state['angle_z'] -= 5
# Distance from origin (+)
elif self.key == ord('H'):
state['radius'] += 0.05
# Distance from origin (-)
elif self.key == ord('G'):
state['radius'] -= 0.05
# Next dataset example
elif self.key == ord(' '):
texture_src = state['dataset'][state['dataset_index']]
state['dataset_index'] += 1
# Next CAD model
elif self.key == ord('N'):
state['cad_idx'] += 1
if state['cad_idx'] == 10:
state['cad_idx'] = 0
# --------------- Update model and Kpoints 3D ----------------------
cad_idx = state['cad_idx']
model_path = args.CAD_root / f'pascal_car_cad_{cad_idx:03d}.ply'
kpoints_3d = state['car_cad_kpoints_3D'][cad_idx].astype(np.float32)
car_mesh = read_triangle_mesh(str(model_path))
# ------------------- Compute normal Colors ----------------------
car_mesh.compute_vertex_normals()
normal_vertex_colors = (np.asarray(car_mesh.vertex_normals) + 1) / 2.
# -------------------- Compute Parts Colors ------------------------
color_dict = pascal_parts_colors['car']
obj_path = model_path.parent / (model_path.stem + '.obj')
color_mesh_from_obj(car_mesh, obj_name=obj_path, mtl_to_color=color_dict,
base_color=color_dict['car'])
part_vertex_colors = np.asarray(car_mesh.vertex_colors).copy()
# Reset colors to normals
car_mesh.vertex_colors = Vector3dVector(normal_vertex_colors)
if 'car_mesh' not in state['geometries']:
state['geometries']['car_mesh'] = car_mesh
else:
state['geometries']['car_mesh'].vertices = car_mesh.vertices
state['geometries']['car_mesh'].vertex_colors = car_mesh.vertex_colors
state['geometries']['car_mesh'].vertex_normals = car_mesh.vertex_normals
state['geometries']['car_mesh'].triangles = car_mesh.triangles
else:
raise NotImplementedError()
# -------------------------- Move Camera ------------------------------
angle_y = np.clip(state['angle_y'], -90, 90)
radius = np.clip(state['radius'], 0, state['radius'])
pascal_az = (state['angle_z'] + 90) % 360
intrinsic = intrinsic_matrix(state['focal'], cx=img_w/2, cy=img_h/2)
if args.verbose:
print(f'Azimuth:{pascal_az} Elevation:{angle_y} Radius:{radius}')
extrinsic = pascal_vpoint_to_extrinsics(az_deg=pascal_az,
el_deg=90 - angle_y,
radius=radius)
# ---------------------- Start Rendering ------------------------------
if not vis.get_render_option() or not vis.get_view_control():
vis.update_geometry() # we don't have anything, return
return
align_view(vis, focal=state['focal'], extrinsic=extrinsic)
vis.get_render_option().mesh_color_option = MeshColorOption.Color
vis.get_render_option().light_on = False
vis.get_render_option().background_color = (0, 0, 0)
# ---------------------- Normal ---------------------------------------
src_shaded = np.asarray(vis.capture_screen_float_buffer(do_render=True))
src_shaded = (src_shaded * 255).astype(np.uint8)
if args.LAB:
src_shaded = cv2.cvtColor(src_shaded, cv2.COLOR_RGB2LAB)
else:
src_shaded = cv2.cvtColor(src_shaded, cv2.COLOR_RGB2BGR)
# ---------------------- Parts ----------------------------------------
state['geometries']['car_mesh'].vertex_colors = Vector3dVector(part_vertex_colors)
vis.update_geometry()
src_parts = np.asarray(vis.capture_screen_float_buffer(do_render=True))
src_parts = (src_parts * 255).astype(np.uint8)
if args.LAB:
src_parts = cv2.cvtColor(src_parts, cv2.COLOR_RGB2LAB)
else:
src_parts = cv2.cvtColor(src_parts, cv2.COLOR_RGB2BGR)
state['geometries']['car_mesh'].vertex_colors = Vector3dVector(normal_vertex_colors)
vis.update_geometry()
# Project model kpoints in 2D
kpoints_2d_step = project_points(kpoints_3d, intrinsic, extrinsic)
kpoints_2d_step /= (img_w, img_h)
kpoints_2d_step = np.clip(kpoints_2d_step, -1, 1)
kpoints_2d_step_dict = {pascal_idx_to_kpoint['car'][i]: kpoints_2d_step[i]
for i in range(kpoints_2d_step.shape[0])}
meta = {'kpoints_2d': kpoints_2d_step_dict,
'vpoint': [pascal_az]}
_, dst_kpoints_planes, dst_visibilities = get_planes(np.zeros((img_h, img_w, 3)),
meta=meta, pascal_class='car')
src_planes = np.asarray([to_image(i, from_LAB=args.LAB) for i in texture_src['planes']])
src_kpoints_planes = texture_src['src_kpoints_planes']
src_visibilities = texture_src['src_vs']
planes_warped, planes_unwarped = warp_unwarp_planes(src_planes=src_planes,
src_planes_kpoints=src_kpoints_planes,
dst_planes_kpoints=dst_kpoints_planes,
src_visibilities=src_visibilities,
dst_visibilities=dst_visibilities,
pascal_class='car')
planes_warped = TextureDatasetWithNormal.planes_to_torch(planes_warped, to_LAB=args.LAB)
planes_warped = planes_warped.reshape(1, planes_warped.shape[0] * planes_warped.shape[1],
planes_warped.shape[2], planes_warped.shape[3])
src_shaded_input = Normalize(mean=[0.5]*3, std=[0.5]*3)(ToTensor()((Image.fromarray(src_shaded))))
src_parts_input = Normalize(mean=[0.5]*3, std=[0.5]*3)(ToTensor()((Image.fromarray(src_parts))))
src_central = texture_src['src_central']
gen_in_src = torch.cat([src_shaded_input.unsqueeze(0), src_parts_input.unsqueeze(0),
src_central.unsqueeze(0), planes_warped], dim=1)
net_image = to_image(state['net'](gen_in_src.to(args.device))[0], from_LAB=args.LAB)
out_image = np.concatenate([to_image(texture_src['src_image'],
from_LAB=args.LAB),
to_image(src_shaded_input, from_LAB=args.LAB),
to_image(src_parts_input, from_LAB=args.LAB),
to_image(src_central, from_LAB=args.LAB),
net_image], axis=1)
state['dump_image'] = out_image
cv2.imshow('Output', out_image)
cv2.waitKey(20)
vis.update_geometry()