blenderEXRdatasetscript.py

import cv2
import bpy
import math
import os
import json
import numpy as np
import subprocess
from mathutils import Vector
import imageio.v3 as iio
import time
import OpenEXR
import Imath

# === CONFIGURATION ===
ASSET_FOLDER = "D:\EXRdataset"
OUTPUT_ROOT = "E:\EXRdatasetout"
IMAGE_RES = (1024, 1024)
VIEWS = [
    (0, 0), (45, 0), (90, 0), (135, 0), (180, 0), (225, 0), (270, 0), (315, 0),
    (0, 45), (90, 45), (180, 45), (270, 45),
    (0, -45), (90, -45), (180, -45), (270, -45),
    (0, 90), (180, 90)
]
EXPOSURE_VALUES = [-3, -2, -1, 0, 1, 2, 3]
MASK_THRESHOLD = 0.5  # For binary mask

def ensure_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def get_object_diameter(obj):
    bbox = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
    min_corner = Vector((min([v[i] for v in bbox]) for i in range(3)))
    max_corner = Vector((max([v[i] for v in bbox]) for i in range(3)))
    diameter = (max_corner - min_corner).length
    return diameter

def get_camera_pose(center, azimuth_deg, elevation_deg, radius):
    az = math.radians(azimuth_deg)
    el = math.radians(elevation_deg)
    x = center[0] + radius * math.cos(el) * math.cos(az)
    y = center[1] + radius * math.cos(el) * math.sin(az)
    z = center[2] + radius * math.sin(el)
    location = (x, y, z)
    direction = Vector(center) - Vector(location)
    rot_quat = direction.to_track_quat('-Z', 'Y')
    return location, rot_quat

def set_camera_intrinsics(camera, focal_length_mm, sensor_width_mm, sensor_height_mm, res_x, res_y):
    camera.data.lens = focal_length_mm
    camera.data.sensor_width = sensor_width_mm
    camera.data.sensor_height = sensor_height_mm
    bpy.context.scene.render.resolution_x = res_x
    bpy.context.scene.render.resolution_y = res_y

def save_metadata(metadata, path):
    with open(path, "w") as f:
        json.dump(metadata, f, indent=2)

def render_passes(output_dir, view_name):
    # Set up render layers for depth, normal, mask
    scene = bpy.context.scene
    scene.use_nodes = True
    tree = scene.node_tree
    links = tree.links
    for node in tree.nodes:
        tree.nodes.remove(node)
    render_layers = tree.nodes.new('CompositorNodeRLayers')
    # Output nodes
    file_output = tree.nodes.new(type="CompositorNodeOutputFile")
    file_output.base_path = output_dir
    file_output.format.file_format = 'OPEN_EXR'
    file_output.file_slots.new("Depth")
    file_output.file_slots.new("Normal")
    file_output.file_slots.new("Mask")
    file_output.file_slots[0].path = f"{view_name}_depth"
    file_output.file_slots[1].path = f"{view_name}_normal"
    file_output.file_slots[2].path = f"{view_name}_mask"
    # Connect outputs
    links.new(render_layers.outputs['Depth'], file_output.inputs[0])
    links.new(render_layers.outputs['Normal'], file_output.inputs[1])
    if 'IndexOB' in render_layers.outputs:
        links.new(render_layers.outputs['IndexOB'], file_output.inputs[2])
    # Render
    bpy.ops.render.render(write_still=True)

def mask_to_svg(mask_path, svg_path):
    print(f"Attempting to vectorize mask: {mask_path} -> {svg_path}")
    # Wait for the mask file to appear (max 5 seconds)
    for _ in range(50):
        if os.path.exists(mask_path):
            break
        time.sleep(0.1)
    else:
        print(f"Mask file not found: {mask_path}")
        return
    mask = iio.imread(mask_path)
    if mask.max() == 0:
        print(f"Mask is empty (all transparent): {mask_path}")
        return

    # Convert to grayscale if needed
    if len(mask.shape) == 3:
        # If RGBA, convert to grayscale using OpenCV
        mask = cv2.cvtColor(mask, cv2.COLOR_RGBA2GRAY)
    elif len(mask.shape) == 2:
        # Already grayscale
        pass
    else:
        print(f"Unexpected mask shape: {mask.shape}")
        return

    # Save as PBM for Potrace
    print("cv2 module:", cv2)
    print("cv2 file:", getattr(cv2, '__file__', 'builtin'))
    print("cv2 has threshold:", hasattr(cv2, "threshold"))
    print("mask shape:", mask.shape)
    _, binary = cv2.threshold(mask, int(MASK_THRESHOLD*255), 255, cv2.THRESH_BINARY)
    temp_pbm = mask_path.replace('.png', '.pbm')
    cv2.imwrite(temp_pbm, binary)
    # Call Potrace
    try:
        result = subprocess.run(["potrace", temp_pbm, "-s", "-o", svg_path], capture_output=True, text=True)
        if result.returncode != 0:
            print(f"Potrace failed: {result.stderr}")
        else:
            print(f"SVG written: {svg_path}")
    except Exception as e:
        print(f"Potrace call failed: {e}")
    # Clean up temp PBM
    if os.path.exists(temp_pbm):
        os.remove(temp_pbm)

def import_asset(asset_path):
    ext = os.path.splitext(asset_path)[1].lower()
    if ext == ".obj":
        bpy.ops.import_scene.obj(filepath=asset_path)
    elif ext == ".fbx":
        bpy.ops.import_scene.fbx(filepath=asset_path)
    elif ext in [".glb", ".gltf"]:
        bpy.ops.import_scene.gltf(filepath=asset_path)
    else:
        raise ValueError(f"Unsupported asset format: {ext}")
    # Return the first mesh object from selected objects
    for obj in bpy.context.selected_objects:
        if obj.type == 'MESH':
            return obj
    raise RuntimeError("No mesh object found after import!")

def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete(use_global=False)

def get_camera_distance_to_fit_object(obj, camera, margin=1.3):
    bbox = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
    min_corner = Vector((min([v[i] for v in bbox]) for i in range(3)))
    max_corner = Vector((max([v[i] for v in bbox]) for i in range(3)))
    size = max_corner - min_corner
    # Use the diagonal in the X/Y plane (width in camera view)
    obj_width = math.sqrt(size.x**2 + size.y**2)
    obj_height = size.z
    aspect = IMAGE_RES[0] / IMAGE_RES[1]
    sensor_fit = camera.data.sensor_fit
    if sensor_fit == 'VERTICAL':
        fov = camera.data.angle_y
        fit_size = obj_height
    else:
        fov = camera.data.angle_x
        fit_size = obj_width
    distance = (fit_size / 2) / math.tan(fov / 2) * margin
    return distance

# Remove exr_alpha_to_png and related EXR-to-PNG alpha extraction logic

# In your main loop, for each view/exposure, before rendering:
# Set up compositor to output alpha as PNG mask
# scene = bpy.context.scene
# scene.use_nodes = True
# tree = scene.node_tree
# tree.nodes.clear()
# rl = tree.nodes.new('CompositorNodeRLayers')
# file_output = tree.nodes.new('CompositorNodeOutputFile')
# file_output.base_path = object_dir
# file_output.format.file_format = 'PNG'
# file_output.file_slots[0].path = f"{view_name}_mask"
# tree.links.new(rl.outputs['Alpha'], file_output.inputs[0])

# Set main render output to EXR
bpy.context.scene.render.image_settings.file_format = 'OPEN_EXR'
bpy.context.scene.render.image_settings.color_depth = '32'
bpy.context.scene.render.image_settings.color_mode = 'RGBA'
# bpy.context.scene.render.filepath = exr_path
bpy.ops.render.render(write_still=True)

# After rendering, the PNG mask will be saved as f"{view_name}_mask0001.png" in object_dir
# mask_path = os.path.join(object_dir, f"{view_name}_mask0001.png")
# mask_to_svg(mask_path, svg_path)

# === MAIN SCRIPT ===

for asset_file in os.listdir(ASSET_FOLDER):
    if not asset_file.lower().endswith(('.obj', '.fbx', '.glb', '.gltf')):
        continue
    clear_scene()
    asset_path = os.path.join(ASSET_FOLDER, asset_file)
    asset_name = os.path.splitext(asset_file)[0]
    object_dir = os.path.join(OUTPUT_ROOT, asset_name)
    ensure_dir(object_dir)
    curves_dir = os.path.join(object_dir, "curves")
    ensure_dir(curves_dir)
    # Set up even lighting (white environment)
    if not bpy.context.scene.world:
        bpy.context.scene.world = bpy.data.worlds.new("World")
    world = bpy.context.scene.world
    world.use_nodes = True
    bg = world.node_tree.nodes.get('Background')
    if bg:
        bg.inputs[0].default_value = (1, 1, 1, 1)  # White
        bg.inputs[1].default_value = 1.0           # Strength
    scene = bpy.context.scene
    scene.view_settings.view_transform = 'Standard'

    # Import asset
    obj = import_asset(asset_path)
    obj.hide_render = False
    obj.hide_set(False)
    obj.pass_index = 1
    print(f"Set obj.pass_index = {obj.pass_index} for {obj.name}")
    # Assign a simple Principled BSDF material with mid-gray color
    if not obj.data.materials:
        mat = bpy.data.materials.new(name="DefaultMat")
        mat.use_nodes = True
        bsdf = mat.node_tree.nodes.get("Principled BSDF")
        if bsdf:
            bsdf.inputs["Base Color"].default_value = (0.8, 0.8, 0.8, 1)
        obj.data.materials.append(mat)
    # Recalculate normals
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.normals_make_consistent(inside=False)
    bpy.ops.object.mode_set(mode='OBJECT')

    # Print all objects, their locations, and types before rendering
    print("Scene objects:")
    for o in bpy.data.objects:
        print(f"  {o.name}: type={o.type}, location={o.location}, scale={o.scale}")

    # Print bounding box size and object scale
    bbox = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
    bbox_min = Vector((min([v[i] for v in bbox]) for i in range(3)))
    bbox_max = Vector((max([v[i] for v in bbox]) for i in range(3)))
    bbox_size = bbox_max - bbox_min
    bbox_center = sum(bbox, Vector((0,0,0))) / 8  # Restored for debug and camera
    print(f"Object bounding box size: {bbox_size}, scale: {obj.scale}")
    if bbox_size.length < 1e-3:
        print("Warning: Object is extremely small!")
    if bbox_size.length > 1e3:
        print("Warning: Object is extremely large!")

    # After importing the object:
    # Create camera if not exists
    if "DatasetCamera" not in bpy.data.objects:
        bpy.ops.object.camera_add()
        camera = bpy.context.active_object
        camera.name = "DatasetCamera"
    else:
        camera = bpy.data.objects["DatasetCamera"]

    # Set camera intrinsics (example values)
    set_camera_intrinsics(camera, 50, 36, 24, IMAGE_RES[0], IMAGE_RES[1])
    bpy.context.scene.camera = camera

    # Compute camera distance to fit object
    CAMERA_DISTANCE = get_camera_distance_to_fit_object(obj, camera, margin=1.3)

    # Enable required passes BEFORE creating the compositor node tree and Render Layers node
    view_layer = bpy.context.view_layer
    view_layer.use_pass_z = True
    view_layer.use_pass_normal = True
    view_layer.use_pass_object_index = True  # Enable before node tree

    # --- Rendering loop ---
    metadata = []
    bpy.context.scene.render.film_transparent = True
    obj.hide_render = False
    obj.hide_set(False)

    # Set frame range to a single frame for animation render
    scene.frame_start = 1
    scene.frame_end = 1
    scene.frame_current = 1

    for idx, (az, el) in enumerate(VIEWS):
        for ev_idx, ev in enumerate(EXPOSURE_VALUES):
            view_name = f"view_{idx:03d}_ev{ev:+d}"
            exr_path = os.path.join(object_dir, f"{view_name}.exr")
            frame = idx * len(EXPOSURE_VALUES) + ev_idx + 1
            scene.frame_current = frame
            scene.update_tag()
            mask_dir = os.path.join(object_dir, "mask")
            ensure_dir(mask_dir)
            mask_path = os.path.join(mask_dir, f"{view_name}_{frame:04d}.png")
            svg_path = os.path.join(curves_dir, f"{view_name}.svg")

            # Re-initialize compositor node tree for each view/EV
            tree = bpy.context.scene.node_tree

            # Set camera pose for this view
            location, rot_quat = get_camera_pose(bbox_center, az, el, CAMERA_DISTANCE)
            camera.location = location
            camera.rotation_mode = 'QUATERNION'
            camera.rotation_quaternion = rot_quat
            print(f"Camera location: {camera.location}, rotation: {camera.rotation_quaternion}")

            # Ensure compositor nodes are enabled before accessing node_tree
            scene = bpy.context.scene
            scene.use_nodes = True
            scene.render.use_compositing = True  # Ensure compositor output is enabled
            tree = scene.node_tree

            # Remove previous output nodes for mask, depth, normal
            for node in tree.nodes:
                if node.type == 'OUTPUT_FILE' and getattr(node, 'label', '') in ['MaskOutput', 'DepthOutput', 'NormalOutput']:
                    tree.nodes.remove(node)

            # Print node tree debug info before rendering
            print('Node tree before rendering:')
            for node in tree.nodes:
                print(f"  Node: {node.name}, type={node.type}, label={getattr(node, 'label', '')}")
                if node.type == 'OUTPUT_FILE':
                    for i, slot in enumerate(node.file_slots):
                        print(f"    Output slot {i}: path={slot.path}, base_path={node.base_path}")

            # Render Layers node (recreate for safety)
            rl = None
            for node in tree.nodes:
                if node.type == 'R_LAYERS':
                    rl = node
                    break
            if rl is None:
                rl = tree.nodes.new('CompositorNodeRLayers')

            # Mask output node
            mask_output = tree.nodes.new('CompositorNodeOutputFile')
            mask_output.label = 'MaskOutput'
            mask_output.base_path = mask_dir
            mask_output.format.file_format = 'PNG'
            mask_output.format.color_depth = '16'
            mask_output.file_slots[0].path = f"{view_name}_####"

            # Depth output node
            depth_output = tree.nodes.new('CompositorNodeOutputFile')
            depth_output.label = 'DepthOutput'
            depth_output.base_path = object_dir
            depth_output.format.file_format = 'OPEN_EXR'
            depth_output.file_slots[0].path = f"{view_name}_depth"

            # Normal output node
            normal_output = tree.nodes.new('CompositorNodeOutputFile')
            normal_output.label = 'NormalOutput'
            normal_output.base_path = object_dir
            normal_output.format.file_format = 'OPEN_EXR'
            normal_output.file_slots[0].path = f"{view_name}_normal"

            # Composite node for main image
            composite = None
            for node in tree.nodes:
                if node.type == 'COMPOSITE':
                    composite = node
                    break
            if composite is None:
                composite = tree.nodes.new('CompositorNodeComposite')
            composite.label = 'MainComposite'

            # Connect outputs
            # Mask: Try ID Mask first, fallback to Alpha
            indexob_output = None
            for o in rl.outputs:
                if o.name == 'IndexOB':
                    indexob_output = o
                    break
            if indexob_output is not None:
                id_mask = tree.nodes.new('CompositorNodeIDMask')
                id_mask.index = 1
                tree.links.new(indexob_output, id_mask.inputs['ID value'])
                tree.links.new(id_mask.outputs['Alpha'], mask_output.inputs[0])
                print("Using ID Mask for PNG mask output.")
            else:
                alpha_output = None
                for o in rl.outputs:
                    if o.name == 'Alpha':
                        alpha_output = o
                        break
                if alpha_output is not None:
                    tree.links.new(alpha_output, mask_output.inputs[0])
                    print("Using Alpha channel for PNG mask output (fallback).")
                else:
                    print("Warning: No Alpha output found! No mask PNG will be created.")
                    mask_output = None

            # Depth
            depth_output_input = None
            for o in rl.outputs:
                if o.name == 'Depth':
                    depth_output_input = o
                    break
            if depth_output_input is not None:
                tree.links.new(depth_output_input, depth_output.inputs[0])
            else:
                print("Warning: No Depth output found!")

            # Normal
            normal_output_input = None
            for o in rl.outputs:
                if o.name == 'Normal':
                    normal_output_input = o
                    break
            if normal_output_input is not None:
                tree.links.new(normal_output_input, normal_output.inputs[0])
            else:
                print("Warning: No Normal output found!")

            # Main image
            image_output = None
            for o in rl.outputs:
                if o.name == 'Image':
                    image_output = o
                    break
            if image_output is not None:
                tree.links.new(image_output, composite.inputs['Image'])
            else:
                print("Warning: No Image output found!")

            bpy.context.scene.update_tag()

            # Set main render output to EXR
            bpy.context.scene.render.image_settings.file_format = 'OPEN_EXR'
            bpy.context.scene.render.image_settings.color_depth = '32'
            bpy.context.scene.render.image_settings.color_mode = 'RGBA'
            bpy.context.scene.render.filepath = exr_path
            bpy.ops.render.render(write_still=True)

            # --- DEBUG: Print node tree structure and connections before rendering ---
            print("[DEBUG] Node tree before render:")
            for node in tree.nodes:
                print(f"  Node: {node.name}, type={node.type}, label={getattr(node, 'label', '')}")
                if node.type == 'OUTPUT_FILE':
                    print(f"    OutputFile base_path: {node.base_path}")
                    for i, slot in enumerate(node.file_slots):
                        print(f"      Slot {i}: path={slot.path}")
                if node.type == 'ID_MASK':
                    print(f"    IDMask index: {node.index}")
            print("[DEBUG] Node links:")
            for link in tree.links:
                print(f"  {link.from_node.name} ({link.from_socket.name}) -> {link.to_node.name} ({link.to_socket.name})")

            # --- After rendering, check if ID Mask output is nonzero ---
            # Only if mask_output is not None
            if mask_output is not None:
                # Try to check if the mask file was written, else fallback to Alpha
                mask_files = os.listdir(mask_dir)
                print("[DEBUG] Files in mask_dir after render:", mask_files)
                found_mask = False
                for fname in mask_files:
                    if fname.startswith(f"{view_name}_") and fname.endswith(".png"):
                        found_mask = True
                        mask_path_actual = os.path.join(mask_dir, fname)
                        print(f"[DEBUG] Found mask file: {mask_path_actual}")
                        mask_to_svg(mask_path_actual, svg_path)
                        break
                if not found_mask:
                    print(f"[DEBUG] Mask file not found for {view_name}, trying Alpha fallback.")
                    # Remove ID Mask node and reconnect Alpha output
                    for node in tree.nodes:
                        if node.type == 'ID_MASK':
                            tree.nodes.remove(node)
                    alpha_output = None
                    for o in rl.outputs:
                        if o.name == 'Alpha':
                            alpha_output = o
                            break
                    if alpha_output is not None:
                        tree.links.new(alpha_output, mask_output.inputs[0])
                        print("[DEBUG] Reconnected Alpha output for mask PNG.")
                        bpy.ops.render.render(write_still=True)
                        mask_files = os.listdir(mask_dir)
                        print("[DEBUG] Files in mask_dir after Alpha fallback render:", mask_files)
                        for fname in mask_files:
                            if fname.startswith(f"{view_name}_") and fname.endswith(".png"):
                                mask_path_actual = os.path.join(mask_dir, fname)
                                print(f"[DEBUG] Found mask file after Alpha fallback: {mask_path_actual}")
                                mask_to_svg(mask_path_actual, svg_path)
                                break
                    else:
                        print("[DEBUG] No Alpha output found for fallback!")

            # Save metadata
            metadata.append({
                "image": f"{view_name}.exr",
                "azimuth": az,
                "elevation": el,
                "ev": ev,
                "camera_distance": CAMERA_DISTANCE,
                "intrinsics": [[camera.data.lens, 0, IMAGE_RES[0]/2], [0, camera.data.lens, IMAGE_RES[1]/2], [0, 0, 1]],
                "extrinsics": [list(camera.matrix_world[i]) for i in range(4)],
                "svg": f"curves/{view_name}.svg",
                "depth": f"{view_name}_depth.exr",
                "normal": f"{view_name}_normal.exr"
            })
    save_metadata(metadata, os.path.join(object_dir, "cameras.json"))
    # Remove asset
    bpy.data.objects.remove(obj, do_unlink=True)

print("Blender dataset generation complete!")