diff --git a/tools/batch_render.sh b/tools/batch_render.sh new file mode 100644 index 0000000..f42c620 --- /dev/null +++ b/tools/batch_render.sh @@ -0,0 +1,5 @@ +#! /bin/bash +python distributed_render.py --num_gpus 8 --workers_per_gpu 2 \ + --data_root /local_home/shenqiuhong/omni3d/raw_scan \ + --output_dir /local_home/shenqiuhong/omni_render/ \ + --limit_num 500 \ No newline at end of file diff --git a/tools/blender_script.py b/tools/blender_script.py new file mode 100644 index 0000000..8850480 --- /dev/null +++ b/tools/blender_script.py @@ -0,0 +1,360 @@ +import sys, os +import json +import bpy +import mathutils +import numpy as np +import pdb +import argparse + + +DEBUG = False + +VIEWS = 100 +RESOLUTION = 512 +#RESULTS_PATH = 'results' +DEPTH_SCALE = 0.1 # 1.4 +COLOR_DEPTH = 8 +FORMAT = 'PNG' +DEPTH_FORMAT = 'OPEN_EXR' +RANDOM_VIEWS = False +UPPER_VIEWS = False # True by default +ANCHOR_VIEWS = True +CIRCLE_FIXED_START = (.3,0,0) +engine = 'BLENDER_EEVEE' + +if ANCHOR_VIEWS: + num_horizontal_views = 8 + elevation_degrees = [0, -15, 15, -30, 30] + elevation_radians = np.radians(elevation_degrees) + azimuth_degrees = np.linspace(0, 2 * np.pi, num_horizontal_views, endpoint=False) + rotation_euler_settings = np.zeros((num_horizontal_views, 3)) + rotation_euler_settings[:, 2] = azimuth_degrees + rotation_euler_settings = rotation_euler_settings[None].repeat(len(elevation_degrees), axis=0) + rotation_euler_settings[:, :, 0] = elevation_radians[:, None].repeat(num_horizontal_views, axis=1) + rotation_euler_settings = rotation_euler_settings.reshape(-1, 3) + VIEWS = num_horizontal_views * len(elevation_degrees) + +def get_scale(obj): + maxWCoord = [0,0,0] + WCoord = [[], [], []] + for vert in obj.data.vertices: + wmtx = obj.matrix_world + worldCoord = vert.co @ wmtx + WCoord[0] += [np.abs(worldCoord[0])] + WCoord[1] += [np.abs(worldCoord[1])] + WCoord[2] += [np.abs(worldCoord[2])] + if len(WCoord[0]) == 0: + return None + maxWCoord[0] = np.max(WCoord[0]) + maxWCoord[1] = np.max(WCoord[1]) + maxWCoord[2] = np.max(WCoord[2]) + max_coord = np.max(maxWCoord) + scale = 0.99 /max_coord + return scale + +def getGeometryCenter(obj): + sumWCoord = [0,0,0] + numbVert = 0 + if obj.type == 'MESH': + for vert in obj.data.vertices: + wmtx = obj.matrix_world + worldCoord = vert.co @ wmtx + sumWCoord[0] += worldCoord[0] + sumWCoord[1] += worldCoord[1] + sumWCoord[2] += worldCoord[2] + numbVert += 1 + sumWCoord[0] = sumWCoord[0]/numbVert + sumWCoord[1] = sumWCoord[1]/numbVert + sumWCoord[2] = sumWCoord[2]/numbVert + return sumWCoord + +def getGeometryBound(obj): + + return sumWCoord + +def setOrigin(obj, r): + oldLoc = obj.location + newLoc = getGeometryCenter(obj) + obj.location.x += (newLoc[0] - oldLoc[0]) * r + obj.location.y += (newLoc[1] - oldLoc[1]) * r + obj.location.z += (newLoc[2] - oldLoc[2]) * r + +def scaleObject(obj, r, adaptive=False): + obj.scale.x *= r + obj.scale.y *= r + obj.scale.z *= r + +##Cleans all decimate modifiers +def cleanAllDecimateModifiers(obj): + for m in obj.modifiers: + if(m.type=="DECIMATE"): + # print("Removing modifier ") + obj.modifiers.remove(modifier=m) + +def render_once(RESULTS_PATH, scale, args): + fp = bpy.path.abspath(f"//{RESULTS_PATH}") + + def listify_matrix(matrix): + matrix_list = [] + for row in matrix: + matrix_list.append(list(row)) + return matrix_list + + if not os.path.exists(fp): + os.makedirs(fp) + + # Data to store in JSON file + out_data = { + 'camera_angle_x': bpy.data.objects['Camera'].data.angle_x, + } + + # Render Optimizations + bpy.context.scene.render.use_persistent_data = True + + context = bpy.context + scene = bpy.context.scene + render = bpy.context.scene.render + + render.engine = engine + # # Set the device_type + # context.preferences.addons[ + # "cycles" + # ].preferences.compute_device_type = "CUDA" # or "OPENCL" + # # Set the device and feature set + # scene.cycles.device = "GPU" + + render.image_settings.color_mode = 'RGBA' # ('RGB', 'RGBA', ...) + render.image_settings.color_depth = str(COLOR_DEPTH) + render.image_settings.file_format = str(FORMAT) + render.resolution_x = RESOLUTION + render.resolution_y = RESOLUTION + render.resolution_percentage = 100 + render.film_transparent = True + + # Set up rendering of depth map. + scene.use_nodes = True + scene.view_layers["View Layer"].use_pass_normal = True + scene.view_layers["View Layer"].use_pass_diffuse_color = True + scene.view_layers["View Layer"].use_pass_object_index = True + + tree = bpy.context.scene.node_tree + nodes = tree.nodes + links = tree.links + # Clear default nodes + for n in nodes: + nodes.remove(n) + + #bpy.context.scene.view_layers["RenderLayer"].use_pass_normal = True + # bpy.context.view_layer.use_pass_z = True + + if not DEBUG: + # Create input render layer node. + render_layers = tree.nodes.new('CompositorNodeRLayers') + + if args.render_depth: + depth_file_output = tree.nodes.new(type="CompositorNodeOutputFile") + depth_file_output.label = 'Depth Output' + depth_file_output.base_path = '' + depth_file_output.file_slots[0].use_node_format = True + depth_file_output.format.file_format = DEPTH_FORMAT + # depth_file_output.format.color_depth = str(COLOR_DEPTH) + if DEPTH_FORMAT == 'OPEN_EXR': + links.new(render_layers.outputs['Depth'], depth_file_output.inputs[0]) + else: + depth_file_output.format.color_mode = "BW" + # Remap as other types can not represent the full range of depth. + map = nodes.new(type="CompositorNodeMapValue") + # Size is chosen kind of arbitrarily, try out until you're satisfied with resulting depth map. + map.offset = [-0.7] + map.size = [DEPTH_SCALE] + map.use_min = True + map.min = [0] + links.new(render_layers.outputs['Depth'], map.inputs[0]) + links.new(map.outputs[0], depth_file_output.inputs[0]) + else: + depth_file_output = None + + if args.render_normal: + normal_file_output = tree.nodes.new(type="CompositorNodeOutputFile") + normal_file_output.label = 'Normal Output' + links.new(render_layers.outputs['Normal'], normal_file_output.inputs[0]) + else: + normal_file_output = None + + # Background + bpy.context.scene.render.dither_intensity = 0.0 + bpy.context.scene.render.film_transparent = True + + # Create collection for objects not to render with background + + objs = [ob for ob in bpy.context.scene.objects if ob.type in ('EMPTY') and 'Empty' in ob.name] + bpy.ops.object.delete({"selected_objects": objs}) + + def parent_obj_to_camera(b_camera): + origin = (0, 0, 0) + b_empty = bpy.data.objects.new("Empty", None) + b_empty.location = origin + b_camera.parent = b_empty # setup parenting + + scn = bpy.context.scene + scn.collection.objects.link(b_empty) + bpy.context.view_layer.objects.active = b_empty + # scn.objects.active = b_empty + return b_empty + + + scene = bpy.context.scene + scene.render.resolution_x = RESOLUTION + scene.render.resolution_y = RESOLUTION + scene.render.resolution_percentage = 100 + + cam = scene.objects['Camera'] + cam.location = (0, 4.0, 0.5) + cam_constraint = cam.constraints.new(type='TRACK_TO') + cam_constraint.track_axis = 'TRACK_NEGATIVE_Z' + cam_constraint.up_axis = 'UP_Y' + b_empty = parent_obj_to_camera(cam) + cam_constraint.target = b_empty + + scene.render.image_settings.file_format = 'PNG' # set output format to .png + + from math import radians + + stepsize = 360.0 / VIEWS + rotation_mode = 'XYZ' + + if not DEBUG: + for output_node in [depth_file_output, normal_file_output]: + if output_node is not None: + output_node.base_path = '' + + out_data['frames'] = [] + + if not RANDOM_VIEWS: + b_empty.rotation_euler = CIRCLE_FIXED_START + + for i in range(0, VIEWS): + if DEBUG: + i = np.random.randint(0,VIEWS) + b_empty.rotation_euler[2] += radians(stepsize*i) + if RANDOM_VIEWS: + scene.render.filepath = fp + '/images/r_' + str(i) + if UPPER_VIEWS: + rot = np.random.uniform(0, 1, size=3) * (1,0,2*np.pi) + rot[0] = np.abs(np.arccos(1 - 2 * rot[0]) - np.pi/2) + b_empty.rotation_euler = rot + else: + b_empty.rotation_euler = np.random.uniform(0, 2*np.pi, size=3) + elif ANCHOR_VIEWS: + scene.render.filepath = fp + '/images/r_' + str(i) + b_empty.rotation_euler = rotation_euler_settings[i] + else: + print("Rotation {}, {}".format((stepsize * i), radians(stepsize * i))) + scene.render.filepath = fp + '/r_' + str(i) + + + if args.render_depth: + depth_file_output.file_slots[0].path = fp + '/depths/r_' + str(i) + "_depth" + if args.render_normal: + normal_file_output.file_slots[0].path = fp + '/normals/r_' + str(i) + "_normal" + + if DEBUG: + break + else: + bpy.ops.render.render(write_still=True) # render still + + frame_data = { + 'file_path': scene.render.filepath.split("/")[-1], # + 'rotation': radians(stepsize), + 'transform_matrix': listify_matrix(cam.matrix_world), + 'scale': scale + } + out_data['frames'].append(frame_data) + + # if RANDOM_VIEWS: + # if UPPER_VIEWS: + # rot = np.random.uniform(0, 1, size=3) * (1,0,2*np.pi) + # rot[0] = np.abs(np.arccos(1 - 2 * rot[0]) - np.pi/2) + # b_empty.rotation_euler = rot + # else: + # b_empty.rotation_euler = np.random.uniform(0, 2*np.pi, size=3) + # else: + # b_empty.rotation_euler[2] += radians(stepsize) + + if not DEBUG: + with open(fp + '/' + 'transforms.json', 'w') as out_file: + json.dump(out_data, out_file, indent=4) + + + +if __name__ == "__main__": + # export DISPLAY=:0.1 && blender --background --python ./blender_script.py -- --obj_path antique/antique_019 + # Make light just directional, disable shadows. + + parser = argparse.ArgumentParser() + parser.add_argument( + "--obj_path", + type=str, + required=True, + help="Path to the object file", + ) + parser.add_argument("--output", type=str, default="./views") + parser.add_argument("--render_normal", action="store_true") + parser.add_argument("--render_depth", action="store_true") + + argv = sys.argv[sys.argv.index("--") + 1 :] + args = parser.parse_args(argv) + + light = bpy.data.lights['Light'] + light.type = 'SUN' + light.use_shadow = False + # Possibly disable specular shading: + light.specular_factor = 0.0 + light.energy = 10.0 + + # Add another light source so stuff facing away from light is not completely dark + bpy.ops.object.light_add(type='SUN') + light2 = bpy.data.lights['Sun'] + light2.use_shadow = False + light2.specular_factor = 0.0 + light2.energy = 1 + bpy.data.objects['Sun'].rotation_euler = bpy.data.objects['Light'].rotation_euler + bpy.data.objects['Sun'].rotation_euler[0] += 180 + + obj_path = args.obj_path + + scan_files = os.listdir(os.path.join(obj_path, 'Scan')) + for scan_file in scan_files: + if '.obj' not in scan_file: + continue + filepath = os.path.join(obj_path, 'Scan', scan_file) + + instance_path = "/".join(obj_path.split("/")[-2:]) + render_path = os.path.join(args.output, instance_path, "render") + + if os.path.exists(render_path): + import shutil + shutil.rmtree(render_path) + bpy.ops.import_scene.obj(filepath=filepath) + scene = bpy.context.scene + mesh_obs = [o for o in scene.objects if o.type == 'MESH'] + + if "Cube" in bpy.data.objects: + bpy.data.objects["Cube"].hide_render = True + bpy.data.objects["Cube"].hide_viewport = True + + if "Cube" in bpy.data.objects: + obj = mesh_obs[1] + else: + obj = mesh_obs[0] + + scale = get_scale(obj) + if scale is None: + bpy.ops.object.delete() + continue + scaleObject(obj, scale) + obj.select_set(True) # Blender 2.8x + + render_once(render_path, scale, args) + bpy.ops.object.delete() + diff --git a/tools/distributed_render.py b/tools/distributed_render.py new file mode 100644 index 0000000..db6d510 --- /dev/null +++ b/tools/distributed_render.py @@ -0,0 +1,110 @@ +import glob +import json +import multiprocessing +import shutil +import subprocess +import time +from dataclasses import dataclass +from typing import Optional + +import tyro +import os +from os import path as osp + +import pdb + + +@dataclass +class Args: + workers_per_gpu: int + """number of workers per gpu""" + + # input_models_path: str + # """Path to a json file containing a list of 3D object files""" + data_root: str + """dataset rootpath""" + + output_dir: str + """output rootpath""" + + num_gpus: int = -1 + """number of gpus to use. -1 means all available gpus""" + + limit_num: int = -1 + """num of objects limit. -1 mean no limit""" + + +def worker( + queue: multiprocessing.JoinableQueue, + count: multiprocessing.Value, + gpu: int, + out_dir: str, +) -> None: + while True: + item = queue.get() + if item is None: + break + + # Perform some operation on the item + print(item, gpu) + command = ( + f"export DISPLAY=:0.{gpu} &&" + f" blender -b -P ./blender_script.py --" + f" --obj_path {item} --output {out_dir}" + ) + subprocess.run(command, shell=True) + + with count.get_lock(): + count.value += 1 + + queue.task_done() + + +if __name__ == "__main__": + # python distributed_render.py --num_gpus 8 --workers_per_gpu 2 --data_root /local_home/shenqiuhong/omni3d/raw_scan --limit_num 500 --output_dir /local_home/shenqiuhong/omni_render/ + args = tyro.cli(Args) + + queue = multiprocessing.JoinableQueue() + count = multiprocessing.Value("i", 0) + + # Start worker processes on each of the GPUs + out_dir = args.output_dir + for gpu_i in range(args.num_gpus): + for worker_i in range(args.workers_per_gpu): + worker_i = gpu_i * args.workers_per_gpu + worker_i + process = multiprocessing.Process( + target=worker, args=(queue, count, gpu_i, out_dir) + ) + process.daemon = True + process.start() + + omni3d_root = args.data_root + + instance_list = [] + + for category in os.listdir(omni3d_root): + category_path = osp.join(omni3d_root, category) + for instance in os.listdir(category_path): + if '.txt' in instance: + instance_path = osp.join(category_path, instance) + os.remove(instance_path) + else: + instance_path = osp.join(category, instance) + instance_list.append(instance_path) + + limit_num = len(instance_list) if args.limit_num < 0 else args.limit_num + instance_list = sorted(instance_list) + instance_list = instance_list[:limit_num] + for instance_path in instance_list: + obj_path = osp.join(omni3d_root, instance_path) + queue.put(obj_path) + + + # Wait for all tasks to be completed + queue.join() + + # Add sentinels to the queue to stop the worker processes + for i in range(args.num_gpus * args.workers_per_gpu): + queue.put(None) + + print("All objects rendered !") diff --git a/tools/single_render.sh b/tools/single_render.sh new file mode 100644 index 0000000..0e3aefe --- /dev/null +++ b/tools/single_render.sh @@ -0,0 +1,4 @@ +#! /bin/bash +export DISPLAY=:0.0 && \ +blender -b -P ./blender_script.py -- --obj_path /local_home/shenqiuhong/omni3d/raw_scan/antique/antique_004 \ +--output /local_home/shenqiuhong/render_test \ \ No newline at end of file diff --git a/tools/start_xserver.py b/tools/start_xserver.py new file mode 100644 index 0000000..d3677c9 --- /dev/null +++ b/tools/start_xserver.py @@ -0,0 +1,271 @@ +# Taken from https://github.com/allenai/ai2thor/blob/main/scripts/ai2thor-xorg +# Starts an x-server to support running Blender on a headless machine with +# dedicated NVIDIA GPUs + +#!/usr/bin/env python3 +import os +import sys +import time +import platform +import re +import shlex +import subprocess +import argparse +import signal + +# Turning off automatic black formatting for this script as it breaks quotes. +# fmt: off +from typing import List + +PID_FILE = "/var/run/ai2thor-xorg.pid" +CONFIG_FILE = "/tmp/ai2thor-xorg.conf" + +DEFAULT_HEIGHT = 768 +DEFAULT_WIDTH = 1024 + + +def process_alive(pid): + """ + Use kill(0) to determine if pid is alive + :param pid: process id + :rtype: bool + """ + try: + os.kill(pid, 0) + except OSError: + return False + + return True + + +def find_devices(excluded_device_ids): + devices = [] + id_counter = 0 + for r in pci_records(): + if r.get("Vendor", "") == "NVIDIA Corporation" and r["Class"] in [ + "VGA compatible controller", + "3D controller", + ]: + bus_id = "PCI:" + ":".join( + map(lambda x: str(int(x, 16)), re.split(r"[:\.]", r["Slot"])) + ) + + if id_counter not in excluded_device_ids: + devices.append(bus_id) + + id_counter += 1 + + if not devices: + print("Error: ai2thor-xorg requires at least one NVIDIA device") + sys.exit(1) + + return devices + +def active_display_bus_ids(): + # this determines whether a monitor is connected to the GPU + # if one is, the following Option is added for the Screen "UseDisplayDevice" "None" + command = "nvidia-smi --query-gpu=pci.bus_id,display_active --format=csv,noheader" + active_bus_ids = set() + result = subprocess.run(command, shell=True, stdout=subprocess.PIPE) + if result.returncode == 0: + for line in result.stdout.decode().strip().split("\n"): + nvidia_bus_id, display_status = re.split(r",\s?", line.strip()) + bus_id = "PCI:" + ":".join( + map(lambda x: str(int(x, 16)), re.split(r"[:\.]", nvidia_bus_id)[1:]) + ) + if display_status.lower() == "enabled": + active_bus_ids.add(bus_id) + + return active_bus_ids + +def pci_records(): + records = [] + command = shlex.split("lspci -vmm") + output = subprocess.check_output(command).decode() + + for devices in output.strip().split("\n\n"): + record = {} + records.append(record) + for row in devices.split("\n"): + key, value = row.split("\t") + record[key.split(":")[0]] = value + + return records + + +def read_pid(): + if os.path.isfile(PID_FILE): + with open(PID_FILE) as f: + return int(f.read()) + else: + return None + + +def start(display: str, excluded_device_ids: List[int], width: int, height: int): + pid = read_pid() + + if pid and process_alive(pid): + print("Error: ai2thor-xorg is already running with pid: %s" % pid) + sys.exit(1) + + with open(CONFIG_FILE, "w") as f: + f.write(generate_xorg_conf(excluded_device_ids, width=width, height=height)) + + log_file = "/var/log/ai2thor-xorg.%s.log" % display + error_log_file = "/var/log/ai2thor-xorg-error.%s.log" % display + command = shlex.split( + "Xorg -quiet -maxclients 1024 -noreset +extension GLX +extension RANDR +extension RENDER -logfile %s -config %s :%s" + % (log_file, CONFIG_FILE, display) + ) + + pid = None + with open(error_log_file, "w") as error_log_f: + proc = subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=error_log_f) + pid = proc.pid + try: + proc.wait(timeout=0.25) + except subprocess.TimeoutExpired: + pass + + if pid and process_alive(pid): + with open(PID_FILE, "w") as f: + f.write(str(proc.pid)) + else: + print("Error: error with command '%s'" % " ".join(command)) + with open(error_log_file, "r") as f: + print(f.read()) + + +def print_config(excluded_device_ids: List[int], width: int, height: int): + print(generate_xorg_conf(excluded_device_ids, width=width, height=height)) + + +def stop(): + pid = read_pid() + if pid and process_alive(pid): + os.kill(pid, signal.SIGTERM) + + for i in range(10): + time.sleep(0.2) + if not process_alive(pid): + os.unlink(PID_FILE) + break + + +def generate_xorg_conf( + excluded_device_ids: List[int], width: int, height: int +): + devices = find_devices(excluded_device_ids) + active_display_devices = active_display_bus_ids() + + xorg_conf = [] + + device_section = """ +Section "Device" + Identifier "Device{device_id}" + Driver "nvidia" + VendorName "NVIDIA Corporation" + BusID "{bus_id}" +EndSection +""" + server_layout_section = """ +Section "ServerLayout" + Identifier "Layout0" + {screen_records} +EndSection +""" + screen_section = """ +Section "Screen" + Identifier "Screen{screen_id}" + Device "Device{device_id}" + DefaultDepth 24 + Option "AllowEmptyInitialConfiguration" "True" + Option "Interactive" "False" + {extra_options} + SubSection "Display" + Depth 24 + Virtual {width} {height} + EndSubSection +EndSection +""" + screen_records = [] + for i, bus_id in enumerate(devices): + extra_options = "" + if bus_id in active_display_devices: + # See https://github.com/allenai/ai2thor/pull/990 + # when a monitor is connected, this option must be used otherwise + # Xorg will fail to start + extra_options = 'Option "UseDisplayDevice" "None"' + xorg_conf.append(device_section.format(device_id=i, bus_id=bus_id)) + xorg_conf.append(screen_section.format(device_id=i, screen_id=i, width=width, height=height, extra_options=extra_options)) + screen_records.append( + 'Screen {screen_id} "Screen{screen_id}" 0 0'.format(screen_id=i) + ) + + xorg_conf.append( + server_layout_section.format(screen_records="\n ".join(screen_records)) + ) + + output = "\n".join(xorg_conf) + return output + + +# fmt: on + +if __name__ == "__main__": + if os.geteuid() != 0: + path = os.path.abspath(__file__) + print("Executing ai2thor-xorg with sudo") + args = ["--", path] + sys.argv[1:] + os.execvp("sudo", args) + + if platform.system() != "Linux": + print("Error: Can only run ai2thor-xorg on linux") + sys.exit(1) + + parser = argparse.ArgumentParser() + parser.add_argument( + "--exclude-device", + help="exclude a specific GPU device", + action="append", + type=int, + default=[], + ) + parser.add_argument( + "--width", + help="width of the screen to start (should be greater than the maximum" + f" width of any ai2thor instance you will start) [default: {DEFAULT_WIDTH}]", + type=int, + default=DEFAULT_WIDTH, + ) + parser.add_argument( + "--height", + help="height of the screen to start (should be greater than the maximum" + f" height of any ai2thor instance you will start) [default: {DEFAULT_HEIGHT}]", + type=int, + default=DEFAULT_HEIGHT, + ) + parser.add_argument( + "command", + help="command to be executed", + choices=["start", "stop", "print-config"], + ) + parser.add_argument( + "display", help="display to be used", nargs="?", type=int, default=0 + ) + args = parser.parse_args() + if args.command == "start": + start( + display=args.display, + excluded_device_ids=args.exclude_device, + height=args.height, + width=args.width, + ) + elif args.command == "stop": + stop() + elif args.command == "print-config": + print_config( + excluded_device_ids=args.exclude_device, + width=args.width, + height=args.height, + )