#!/usr/bin/env python3
import _init_paths
import time
import os
import contextlib
from math import radians
import numpy as np
from PIL import Image
from tempfile import TemporaryFile
from lib.utils import stdout_redirected
from lib.config import cfg
import bpy
def voxel2mesh(voxels):
cube_verts = [[0, 0, 0],
[0, 0, 1],
[0, 1, 0],
[0, 1, 1],
[1, 0, 0],
[1, 0, 1],
[1, 1, 0],
[1, 1, 1]] # 8 points
cube_faces = [[0, 1, 2],
[1, 3, 2],
[2, 3, 6],
[3, 7, 6],
[0, 2, 6],
[0, 6, 4],
[0, 5, 1],
[0, 4, 5],
[6, 7, 5],
[6, 5, 4],
[1, 7, 3],
[1, 5, 7]] # 12 face
cube_verts = np.array(cube_verts)
cube_faces = np.array(cube_faces) + 1
l, m, n = voxels.shape
scale = 0.01
cube_dist_scale = 1.1
verts = []
faces = []
curr_vert = 0
for i in range(l):
for j in range(m):
for k in range(n):
# If there is a non-empty voxel
if voxels[i, j, k] > 0:
verts.extend(scale * (cube_verts + cube_dist_scale * np.array([[i, j, k]])))
faces.extend(cube_faces + curr_vert)
curr_vert += len(cube_verts)
return np.array(verts), np.array(faces)
def write_obj(filename, verts, faces):
""" write the verts and faces on file."""
with open(filename, 'w') as f:
# write vertices
f.write('g\n# %d vertex\n' % len(verts))
for vert in verts:
f.write('v %f %f %f\n' % tuple(vert))
# write faces
f.write('# %d faces\n' % len(faces))
for face in faces:
f.write('f %d %d %d\n' % tuple(face))
class BaseRenderer:
model_idx = 0
def __init__(self):
# bpy.data.scenes['Scene'].render.engine = 'CYCLES'
# bpy.context.scene.cycles.device = 'GPU'
# bpy.context.user_preferences.system.compute_device_type = 'CUDA'
# bpy.context.user_preferences.system.compute_device = 'CUDA_1'
# changing these values does affect the render.
# remove the default cube
bpy.ops.object.select_pattern(pattern="Cube")
bpy.ops.object.delete()
render_context = bpy.context.scene.render
world = bpy.context.scene.world
camera = bpy.data.objects['Camera']
light_1 = bpy.data.objects['Lamp']
light_1.data.type = 'HEMI'
# set the camera postion and orientation so that it is in
# the front of the object
camera.location = (1, 0, 0)
camera.rotation_mode = 'ZXY'
camera.rotation_euler = (0, radians(90), radians(90))
# parent camera with a empty object at origin
org_obj = bpy.data.objects.new("RotCenter", None)
org_obj.location = (0, 0, 0)
org_obj.rotation_euler = (0, 0, 0)
bpy.context.scene.objects.link(org_obj)
camera.parent = org_obj # setup parenting
# render setting
render_context.resolution_percentage = 100
world.horizon_color = (1, 1, 1) # set background color to be white
# set file name for storing rendering result
self.result_fn = '%s/render_result_%d.png' % (cfg.DIR.RENDERING_PATH, os.getpid())
bpy.context.scene.render.filepath = self.result_fn
self.render_context = render_context
self.org_obj = org_obj
self.camera = camera
self.light = light_1
self._set_lighting()
def initialize(self, models_fn, viewport_size_x, viewport_size_y):
self.models_fn = models_fn
self.render_context.resolution_x = viewport_size_x
self.render_context.resolution_y = viewport_size_y
def _set_lighting(self):
pass
def setViewpoint(self, azimuth, altitude, yaw, distance_ratio, fov):
self.org_obj.rotation_euler = (0, 0, 0)
self.light.location = (distance_ratio *
(cfg.RENDERING.MAX_CAMERA_DIST + 2), 0, 0)
self.camera.location = (distance_ratio *
cfg.RENDERING.MAX_CAMERA_DIST, 0, 0)
self.org_obj.rotation_euler = (radians(-yaw),
radians(-altitude),
radians(-azimuth))
def setTransparency(self, transparency):
""" transparency is either 'SKY', 'TRANSPARENT'
If set 'SKY', render background using sky color."""
self.render_context.alpha_mode = transparency
def selectModel(self):
bpy.ops.object.select_all(action='DESELECT')
bpy.ops.object.select_pattern(pattern="RotCenter")
bpy.ops.object.select_pattern(pattern="Lamp*")
bpy.ops.object.select_pattern(pattern="Camera")
bpy.ops.object.select_all(action='INVERT')
def printSelection(self):
print(bpy.context.selected_objects)
def clearModel(self):
self.selectModel()
bpy.ops.object.delete()
# The meshes still present after delete
for item in bpy.data.meshes:
bpy.data.meshes.remove(item)
for item in bpy.data.materials:
bpy.data.materials.remove(item)
def setModelIndex(self, model_idx):
self.model_idx = model_idx
def loadModel(self, file_path=None):
if file_path is None:
file_path = self.models_fn[self.model_idx]
if file_path.endswith('obj'):
bpy.ops.import_scene.obj(filepath=file_path)
elif file_path.endswith('3ds'):
bpy.ops.import_scene.autodesk_3ds(filepath=file_path)
elif file_path.endswith('dae'):
# Must install OpenCollada. Please read README.md
bpy.ops.wm.collada_import(filepath=file_path)
else:
raise Exception("Loading failed: %s Model loading for type %s not Implemented" %
(file_path, file_path[-4:]))
def render(self, load_model=True, clear_model=True, resize_ratio=None,
return_image=True, image_path=os.path.join(cfg.RENDERING.BLENDER_TMP_DIR, 'tmp.png')):
""" Render the object """
if load_model:
self.loadModel()
# resize object
self.selectModel()
if resize_ratio:
bpy.ops.transform.resize(value=resize_ratio)
self.result_fn = image_path
bpy.context.scene.render.filepath = image_path
bpy.ops.render.render(write_still=True) # save straight to file
if resize_ratio:
bpy.ops.transform.resize(value=(1/resize_ratio[0],
1/resize_ratio[1], 1/resize_ratio[2]))
if clear_model:
self.clearModel()
if return_image:
im = np.array(Image.open(self.result_fn)) # read the image
# Last channel is the alpha channel (transparency)
return im[:, :, :3], im[:, :, 3]
class ShapeNetRenderer(BaseRenderer):
def __init__(self):
super().__init__()
self.setTransparency('TRANSPARENT')
def _set_lighting(self):
# Create new lamp datablock
light_data = bpy.data.lamps.new(name="New Lamp", type='HEMI')
# Create new object with our lamp datablock
light_2 = bpy.data.objects.new(name="New Lamp", object_data=light_data)
bpy.context.scene.objects.link(light_2)
# put the light behind the camera. Reduce specular lighting
self.light.location = (0, -2, 2)
self.light.rotation_mode = 'ZXY'
self.light.rotation_euler = (radians(45), 0, radians(90))
self.light.data.energy = 0.7
light_2.location = (0, 2, 2)
light_2.rotation_mode = 'ZXY'
light_2.rotation_euler = (-radians(45), 0, radians(90))
light_2.data.energy = 0.7
class VoxelRenderer(BaseRenderer):
def __init__(self):
super().__init__()
self.setTransparency('SKY')
def _set_lighting(self):
self.light.location = (0, 3, 3)
self.light.rotation_mode = 'ZXY'
self.light.rotation_euler = (-radians(45), 0, radians(90))
self.light.data.energy = 0.7
# Create new lamp datablock
light_data = bpy.data.lamps.new(name="New Lamp", type='HEMI')
# Create new object with our lamp datablock
light_2 = bpy.data.objects.new(name="New Lamp", object_data=light_data)
bpy.context.scene.objects.link(light_2)
light_2.location = (4, 1, 6)
light_2.rotation_mode = 'XYZ'
light_2.rotation_euler = (radians(37), radians(3), radians(106))
light_2.data.energy = 0.7
def render_voxel(self, pred, thresh=0.4,
image_path=os.path.join(cfg.RENDERING.BLENDER_TMP_DIR, 'tmp.png')):
# Cleanup the scene
self.clearModel()
out_f = os.path.join(cfg.RENDERING.BLENDER_TMP_DIR, 'tmp.obj')
occupancy = pred > thresh
vertices, faces = voxel2mesh(occupancy)
with contextlib.suppress(IOError):
os.remove(out_f)
write_obj(out_f, vertices, faces)
# Load the obj
bpy.ops.import_scene.obj(filepath=out_f)
bpy.context.scene.render.filepath = image_path
bpy.ops.render.render(write_still=True) # save straight to file
im = np.array(Image.open(image_path)) # read the image
# Last channel is the alpha channel (transparency)
return im[:, :, :3], im[:, :, 3]
def main():
"""Test function"""
# Modify the following file to visualize the model
dn = '/ShapeNet/ShapeNetCore.v1/02958343/'
model_id = [line.strip('\n') for line in open(dn + 'models.txt')]
file_paths = [os.path.join(dn, line, 'model.obj') for line in model_id]
sum_time = 0
renderer = ShapeNetRenderer()
renderer.initialize(file_paths, 500, 500)
for i, curr_model_id in enumerate(model_id):
start = time.time()
image_path = '%s/%s.png' % ('/tmp', curr_model_id[:-4])
az, el, depth_ratio = list(
*([360, 5, 0.3] * np.random.rand(1, 3) + [0, 25, 0.65]))
renderer.setModelIndex(i)
renderer.setViewpoint(30, 30, 0, 0.7, 25)
with TemporaryFile() as f, stdout_redirected(f):
rendering, alpha = renderer.render(load_model=True,
clear_model=True, image_path=image_path)
print('Saved at %s' % image_path)
end = time.time()
sum_time += end - start
if i % 10 == 0:
print(sum_time/(10))
sum_time = 0
if __name__ == "__main__":
main()
