roi/EditP23
5
1# This script is borrowed from https://github.com/allenai/objaverse-rendering2import argparse3import math4import os5from pathlib import Path6import shutil7from typing import Dict, Literal, Tuple8 9import bpy10from mathutils import Vector11from PIL import Image12 13 14 15# --- Blender Setup Functions ---16def global_settings():17 """Configures global Blender rendering settings."""18 context = bpy.context19 scene = context.scene20 render = scene.render21 22 render.engine = "CYCLES"23 render.image_settings.file_format = "PNG"24 render.image_settings.color_mode = "RGBA"25 render.resolution_x = 51226 render.resolution_y = 51227 render.resolution_percentage = 10028 29 scene.cycles.device = "GPU"30 scene.cycles.samples = 3231 scene.cycles.diffuse_bounces = 132 scene.cycles.glossy_bounces = 133 scene.cycles.transparent_max_bounces = 334 scene.cycles.transmission_bounces = 335 scene.cycles.filter_width = 0.0136 scene.cycles.use_denoising = True37 scene.render.film_transparent = True38 return scene39 40 41def add_lighting() -> None:42 """Adds area lights to the scene."""43 # Delete the default light44 if "Light" in bpy.data.objects:45 bpy.data.objects["Light"].select_set(True)46 bpy.ops.object.delete()47 48 # Add a new large area light49 bpy.ops.object.light_add(type="AREA")50 light2 = bpy.data.lights["Area"]51 light2.energy = 3000052 bpy.data.objects["Area"].location[2] = 0.553 bpy.data.objects["Area"].scale[0] = 10054 bpy.data.objects["Area"].scale[1] = 10055 bpy.data.objects["Area"].scale[2] = 10056 57 # Add a fill light58 bpy.ops.object.light_add(type="AREA", location=(0, 0, 2))59 fill_obj = bpy.context.active_object60 fill_obj.data.energy = 200061 fill_obj.scale = (10, 10, 10)62 63 64def reset_scene() -> None:65 """Resets the scene to a clean state by deleting all objects and data."""66 # Delete all objects67 bpy.ops.object.select_all(action='SELECT')68 bpy.ops.object.delete()69 70 # Delete all meshes71 for block in bpy.data.meshes:72 bpy.data.meshes.remove(block, do_unlink=True)73 74 # Delete all materials75 for material in bpy.data.materials:76 bpy.data.materials.remove(material, do_unlink=True)77 78 # Delete all textures79 for texture in bpy.data.textures:80 bpy.data.textures.remove(texture, do_unlink=True)81 82 # Delete all images83 for image in bpy.data.images:84 bpy.data.images.remove(image, do_unlink=True)85 86 # Delete all lights87 for light in bpy.data.lights:88 bpy.data.lights.remove(light, do_unlink=True)89 90 # Delete all cameras91 for cam in bpy.data.cameras:92 bpy.data.cameras.remove(cam, do_unlink=True)93 94 # Delete all empties and curves95 for curve in bpy.data.curves:96 bpy.data.curves.remove(curve, do_unlink=True)97 98 # Reset world99 if bpy.data.worlds:100 for world in bpy.data.worlds:101 bpy.data.worlds.remove(world, do_unlink=True)102 103 # Create a new default world104 bpy.context.scene.world = bpy.data.worlds.new("World")105 bpy.context.view_layer.update()106 107 108def load_object(object_path: str) -> None:109 """Loads a 3D model into the scene based on its file extension."""110 if object_path.endswith(".glb"):111 bpy.ops.import_scene.gltf(filepath=object_path, merge_vertices=True)112 elif object_path.endswith(".fbx"):113 bpy.ops.import_scene.fbx(filepath=object_path)114 else:115 raise ValueError(f"Unsupported file type: {object_path}")116 117 118# --- Scene Normalization and Utility Functions ---119def scene_bbox(single_obj=None, ignore_matrix=False):120 """Calculates the bounding box of the scene or a single object."""121 bbox_min = (math.inf,) * 3122 bbox_max = (-math.inf,) * 3123 found = False124 for obj in scene_meshes() if single_obj is None else [single_obj]:125 found = True126 for coord in obj.bound_box:127 coord = Vector(coord)128 if not ignore_matrix:129 coord = obj.matrix_world @ coord130 bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))131 bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))132 if not found:133 raise RuntimeError("No objects in scene to compute bounding box for")134 return Vector(bbox_min), Vector(bbox_max)135 136 137def scene_root_objects():138 """Generator for all root objects in the scene."""139 for obj in bpy.context.scene.objects.values():140 if not obj.parent:141 yield obj142 143 144def scene_meshes():145 """Generator for all mesh objects in the scene."""146 for obj in bpy.context.scene.objects.values():147 if isinstance(obj.data, (bpy.types.Mesh)):148 yield obj149 150 151def normalize_scene(target_scale=1.0):152 """Normalizes the scene: scales to fit target size and centers at the origin."""153 bbox_min, bbox_max = scene_bbox()154 size = bbox_max - bbox_min155 max_dim = max(size.x, size.y, size.z)156 if max_dim == 0:157 raise ValueError("Model has zero size. Cannot normalize.")158 159 scale = target_scale / max_dim160 for obj in scene_root_objects():161 obj.scale = obj.scale * scale162 163 bpy.context.view_layer.update()164 165 bbox_min, bbox_max = scene_bbox()166 center = (bbox_min + bbox_max) * 0.5167 for obj in scene_root_objects():168 obj.location -= center169 170 bpy.context.view_layer.update()171 172 173# --- Camera and Lighting Setup ---174def setup_camera(scene):175 """Configures the camera and adds a tracking constraint."""176 cam = scene.objects["Camera"]177 cam.location = (0, 1.2, 0)178 cam.data.lens = 35179 cam.data.sensor_width = 32180 cam_constraint = cam.constraints.new(type="TRACK_TO")181 cam_constraint.track_axis = "TRACK_NEGATIVE_Z"182 cam_constraint.up_axis = "UP_Y"183 return cam, cam_constraint184 185 186def _create_light(187 name: str,188 light_type: Literal["POINT", "SUN", "SPOT", "AREA"],189 location: Tuple[float, float, float],190 rotation: Tuple[float, float, float],191 energy: float,192 use_shadow: bool = False,193 specular_factor: float = 1.0,194) -> bpy.types.Object:195 """Creates and returns a configured light object."""196 light_data = bpy.data.lights.new(name=name, type=light_type)197 light_object = bpy.data.objects.new(name, light_data)198 bpy.context.collection.objects.link(light_object)199 200 light_object.location = location201 light_object.rotation_euler = rotation202 203 light_data.energy = energy204 light_data.use_shadow = use_shadow205 light_data.specular_factor = specular_factor206 207 return light_object208 209 210def create_lighting() -> Dict[str, bpy.types.Object]:211 """Creates a deterministic multi-directional sun lighting setup."""212 # Remove existing lights213 bpy.ops.object.select_all(action="DESELECT")214 bpy.ops.object.select_by_type(type="LIGHT")215 bpy.ops.object.delete()216 217 # Add 4 deterministic sun lights218 key_light = _create_light(219 name="Key_Light",220 light_type="SUN",221 location=(0, 0, 0),222 rotation=(0.785398, 0, -0.785398), # 45°, -45° in radians223 energy=0.5,224 )225 fill_light = _create_light(226 name="Fill_Light",227 light_type="SUN",228 location=(0, 0, 0),229 rotation=(0.785398, 0, 2.35619), # 45°, 135°230 energy=0.3,231 )232 rim_light = _create_light(233 name="Rim_Light",234 light_type="SUN",235 location=(0, 0, 0),236 rotation=(-0.785398, 0, -3.92699), # -45°, -225°237 energy=0.5,238 )239 bottom_light = _create_light(240 name="Bottom_Light",241 light_type="SUN",242 location=(0, 0, 0),243 rotation=(3.14159, 0, 0), # 180° (from below)244 energy=0.2,245 )246 return {247 "key_light": key_light,248 "fill_light": fill_light,249 "rim_light": rim_light,250 "bottom_light": bottom_light,251 }252 253 254# --- Main Rendering and Image Processing Functions ---255def render_object(256 object_file: str,257 output_dir: str,258 camera_views=[(30, 30, 1.5), (90, -20, 1.5), (150, 30, 1.5), (210, -20, 1.5), (270, 30, 1.5), (330, -20, 1.5)],259 background_color=(255, 255, 255)260) -> None:261 """Renders images of an object from multiple camera views."""262 scene = global_settings()263 os.makedirs(output_dir, exist_ok=True)264 reset_scene()265 266 # Create and set up a new camera267 bpy.ops.object.camera_add()268 camera = bpy.context.object269 camera.name = "Camera"270 scene.collection.objects.link(camera)271 scene.camera = camera272 273 scene.view_settings.view_transform = 'Standard'274 275 # Set background color276 world = bpy.data.worlds["World"]277 world.use_nodes = False278 world.color = tuple(channel / 255 for channel in background_color)279 scene.render.film_transparent = False280 scene.world = world281 282 # Load, normalize, and light the object283 load_object(object_file)284 normalize_scene()285 create_lighting()286 cam, cam_constraint = setup_camera(scene)287 288 # Create an empty object for the camera to track289 empty = bpy.data.objects.new("Empty", None)290 scene.collection.objects.link(empty)291 cam_constraint.target = empty292 293 for i, (azim, elev, camera_dist) in enumerate(camera_views):294 # Set camera position295 theta = math.radians(azim)296 phi = math.radians(elev)297 point = (298 camera_dist * math.cos(phi) * math.cos(theta),299 camera_dist * math.cos(phi) * math.sin(theta),300 camera_dist * math.sin(phi),301 )302 cam.location = point303 304 # Render the image305 render_path = os.path.join(output_dir, f"{i:02d}.png")306 scene.render.filepath = render_path307 bpy.ops.render.render(write_still=True)308 309 310def create_tiled_grid(311 image_paths=["00.png", "01.png", "02.png", "03.png", "04.png", "05.png"],312 output_path="tiled_grid.png",313 tile_width=320,314 tile_height=320,315 background_color=(255, 255, 255),316):317 """Creates a 2x3 tiled grid image from a list of six image paths."""318 if len(image_paths) != 6:319 print("Error: Exactly 6 image paths are required.")320 return321 322 grid_width = tile_width * 2323 grid_height = tile_height * 3324 grid_image = Image.new("RGB", (grid_width, grid_height), background_color)325 326 for i, image_path in enumerate(image_paths):327 img = Image.open(image_path)328 img = img.resize((tile_width, tile_height))329 # Handle transparency by pasting onto a solid background330 if img.mode == "RGBA":331 background = Image.new("RGB", (tile_width, tile_height), background_color)332 background.paste(img, (0, 0), img)333 img = background334 335 x = (i % 2) * tile_width336 y = (i // 2) * tile_height337 grid_image.paste(img, (x, y))338 339 grid_image.save(output_path)340 print(f"Tiled grid image saved to: {output_path}")341 342 343 344# --- Main Execution Block ---345if __name__ == "__main__":346 parser = argparse.ArgumentParser(description="Render a 3D object into a multi-view and source image format for EditP23.")347 parser.add_argument("--mesh_path", type=str, required=True, help="Path to the input .glb or .fbx file.")348 parser.add_argument("--output_dir", type=str, required=True, help="Directory to save the output src.png and src_mv.png.")349 parser.add_argument("--camera_dist", type=float, default=1.35, help="Camera distance from the object.")350 parser.add_argument("--azim_offset", type=float, default=0, help="Azimuthal offset for camera views in degrees.")351 args = parser.parse_args()352 353 RENDERS_SUBDIR = "all_renders"354 BACKGROUND_COLOR = (255, 255, 255)355 356 output_dir = Path(args.output_dir)357 renders_path = output_dir / RENDERS_SUBDIR358 359 360 ELEV_1 = 20361 ELEV_2 = -10362 elevs = [ELEV_1, ELEV_2] * 3363 azims = [(30 + 60 * i + args.azim_offset) % 360 for i in range(6)]364 camera_views = [(azim, elev, args.camera_dist) for azim, elev in zip(azims, elevs)] + [365 ((0 + args.azim_offset) % 360, ELEV_1, args.camera_dist)366 ]367 368 369 # Render the object from different views370 render_object(371 args.mesh_path,372 output_dir=str(renders_path),373 camera_views=camera_views,374 background_color=BACKGROUND_COLOR,375 )376 377 # --- Create Final Outputs ---378 image_paths_for_grid = [renders_path / f"{i:02d}.png" for i in range(6)]379 380 create_tiled_grid(381 image_paths=image_paths_for_grid,382 output_path=str(output_dir/"src_mv.png"),383 background_color=BACKGROUND_COLOR,384 )385 386 shutil.copy(renders_path / "06.png", output_dir / "src.png")387 388 print(f"Saved conditioning view and multi-view grid to {renders_path}.")389