nodes_xyz/scripts/mark_scene_entities_as_asset.py

#!/usr/bin/python3

import os
import pathlib
import random
import math
from posixpath import relpath
import tempfile

import bpy

# Config

tags=[""]
force_preview_update = True

# Nothing to change below

def find_catalog_id():
    pathsUUIDs = {}
    blender_catalogs_file = "blender_assets.cats.txt"
    current_dir = pathlib.Path().absolute()
    blender_catalogs_dir = None
    catalog_id = None

    for x in [current_dir, *pathlib.Path(current_dir).parents]:
        if os.path.exists(os.path.join(x, blender_catalogs_file)):
            blender_catalogs_dir = x
            break

    if blender_catalogs_dir is not None:
        with open(os.path.join(blender_catalogs_dir, blender_catalogs_file)) as f:
            for line in f.readlines():
                if not line.startswith("#") and ":" in line:
                    line_uuid = line.split(":")[0].strip()
                    path = line.split(":")[1].strip()
                    pathsUUIDs[path] = line_uuid
        rel_path = os.path.relpath(current_dir, blender_catalogs_dir)
        while rel_path != ".":
            if rel_path in pathsUUIDs:
                catalog_id = pathsUUIDs[rel_path]
                break
            rel_path = str(pathlib.Path(rel_path).parent)
    return catalog_id


def setup_world(scene):
    # Make sure we have a camera
    camera_data = bpy.data.cameras.new(name='Camera')
    camera = bpy.data.objects.new('Camera', camera_data)
    scene.camera = camera
    bpy.context.scene.collection.objects.link(camera)
        
    # Change Settings
    camera.rotation_euler = (70/180*math.pi, 0, -20/180*math.pi)
    # This was needed for very small assets. We could base the clip planes on the scene's bounding box.
    # camera_data.clip_start = 0.001
    # camera_data.clip_end = 10
    scene.render.engine = 'CYCLES'
    scene.render.resolution_y = 256
    scene.render.resolution_x = 256
    scene.render.film_transparent = True
    scene.render.image_settings.file_format = 'PNG'

    # Setup the environment map
    scene.world.use_nodes = True
    world_tree = scene.world.node_tree
    env_node = world_tree.nodes.new('ShaderNodeTexEnvironment')
    bg_node = world_tree.nodes['Background']
    world_tree.links.new(bg_node.inputs['Color'], env_node.outputs['Color'])
    env_node.image = bpy.data.images.load("/windows/d/3DLibrary/hdri/HDRIHaven/old_depot_4k.hdr")


def snapshot(scene, entity, tmpdirname):
    bpy.ops.object.select_all(action='DESELECT')
    if entity.rna_type.name == 'Collection':
        for o in entity.objects:
            o.select_set(True)
    else:
        entity.select_set(True)
    bpy.ops.view3d.camera_to_view_selected()
        
    filename = str(random.randint(0,100000000000))+".png"
    filepath = str(os.path.abspath(os.path.join(tmpdirname, filename)))
    scene.render.filepath = filepath
    
    #Render File, Mark Asset and Set Image
    bpy.ops.render.render(write_still = True)
    
    return filepath


def ensure_preview_image(obj, tmpdirname):
    preview_filepath = None
    if obj.preview is None or force_preview_update:
        exts=['.jpg', '.jpeg', '.png']
        for ext in exts:
            check_filepath = bpy.data.filepath.replace(".blend", ext)
            if os.path.exists(check_filepath):
                preview_filepath = check_filepath
        if preview_filepath is None:
            print("Generating preview image")
            preview_filepath = snapshot(bpy.context.scene, obj, tmpdirname)
    return preview_filepath

def mark_entity(obj, preview_filepath, catalog_id):
    # Mark asset
    if obj.asset_data is None:
        print("Marking ", obj)
        obj.asset_mark()
    # Set tags
    for tag in tags:
        if not tag in obj.asset_data.tags:
            print("Creating tag ", tag)
            obj.asset_data.tags.new(tag)
    # Set preview
    if obj.preview is None or force_preview_update:
        override = bpy.context.copy()
        override['id'] = obj
        print("Loading preview image: ", preview_filepath)
        with bpy.context.temp_override(**override):
            bpy.ops.ed.lib_id_load_custom_preview(filepath=preview_filepath)
    # Put into right catalog
    if catalog_id is not None:
        print("Setting catalog id:", catalog_id)
        obj.asset_data.catalog_id = catalog_id


def collect_asset_entities():
    asset_entities = []
    if len(bpy.data.objects)==1:
        for obj in bpy.data.objects:    
            if obj.parent is None:
                asset_entities.append(obj)
    else:
        # We are only interested in root collections, which we can find as
        # children of the scene collection
        for col in bpy.context.scene.collection.children:
            asset_entities.append(col)
    return asset_entities


with tempfile.TemporaryDirectory() as tmpdirname:
    setup_world(bpy.context.scene)
    preview_file_paths = []
    for entity in collect_asset_entities():
        preview_file_paths.append(ensure_preview_image(entity, tmpdirname))
    #Cleanup
    bpy.ops.wm.revert_mainfile()
    # We need to collect the entities twice, because we might revert the blend file inbetween
    for entity,preview_file_path in zip(collect_asset_entities(), preview_file_paths):
        mark_entity(entity, preview_file_path, find_catalog_id())
    bpy.ops.wm.save_mainfile()