core/assets/basket.py

# Copyright (C) 2023, Princeton University.
# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.

# Authors: Beining Han

import bpy
import numpy as np
from numpy.random import uniform
import random
import time

from infinigen.assets.materials.plastics.plastic_rough import shader_rough_plastic
from infinigen.core import surface, tagging
from infinigen.core.nodes import node_utils
from infinigen.core.nodes.node_wrangler import Nodes, NodeWrangler
from infinigen.core.placement.factory import AssetFactory


@node_utils.to_nodegroup("nodegroup_holes", singleton=False, type="GeometryNodeTree")
def nodegroup_holes(nw: NodeWrangler):
    # Code generated using version 2.6.4 of the node_transpiler

    group_input = nw.new_node(
        Nodes.GroupInput,
        expose_input=[
            ("NodeSocketFloat", "height", 0.5000),
            ("NodeSocketFloat", "gap_size", 0.5000),
            ("NodeSocketFloat", "hole_edge_gap", 0.5000),
            ("NodeSocketFloat", "hole_size", 0.5000),
            ("NodeSocketFloat", "depth", 0.5000),
            ("NodeSocketFloat", "width", 0.5000),
        ],
    )

    add = nw.new_node(
        Nodes.Math,
        input_kwargs={0: group_input.outputs["hole_edge_gap"], 1: 0.0000},
        attrs={"operation": "ADD"}
    )

    subtract = nw.new_node(
        Nodes.Math,
        input_kwargs={0: group_input.outputs["height"], 1: add},
        attrs={"operation": "SUBTRACT"},
    )

    add_1 = nw.new_node(
        Nodes.Math, 
        input_kwargs={0: group_input.outputs["width"], 1: 0.0000},
        attrs={"operation": "ADD"}
    )

    subtract_1 = nw.new_node(
        Nodes.Math, input_kwargs={0: add_1, 1: add}, attrs={"operation": "SUBTRACT"}
    )

    add_2 = nw.new_node(
        Nodes.Math, input_kwargs={0: group_input.outputs["hole_size"], 1: 0.0000}, attrs={"operation": "ADD"}
    )

    add_3 = nw.new_node(
        Nodes.Math, input_kwargs={0: add_2, 1: group_input.outputs["gap_size"]}, attrs={"operation": "ADD"}
    )

    divide = nw.new_node(
        Nodes.Math, input_kwargs={0: subtract, 1: add_3}, attrs={"operation": "DIVIDE"}
    )

    divide_1 = nw.new_node(
        Nodes.Math,
        input_kwargs={0: subtract_1, 1: add_3},
        attrs={"operation": "DIVIDE"},
    )

    grid = nw.new_node(
        Nodes.MeshGrid,
        input_kwargs={
            "Size X": subtract,
            "Size Y": subtract_1,
            "Vertices X": divide,
            "Vertices Y": divide_1,
        },
    )

    store_named_attribute = nw.new_node(
        Nodes.StoreNamedAttribute,
        input_kwargs={
            "Geometry": grid.outputs["Mesh"],
            "Name": "uv_map",
            3: grid.outputs["UV Map"],
        },
        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
    )

    transform_1 = nw.new_node(
        Nodes.Transform,
        input_kwargs={
            "Geometry": store_named_attribute,
            "Rotation": (0.0000, 1.5708, 0.0000),
        },
    )

    add_4 = nw.new_node(
        Nodes.Math, input_kwargs={0: group_input.outputs["depth"], 1: 0.0000}, attrs={"operation": "ADD"}
    )

    add_5 = nw.new_node(Nodes.Math, input_kwargs={0: add_4, 1: 0.1}, attrs={"operation": "ADD"})

    combine_xyz_3 = nw.new_node(
        Nodes.CombineXYZ, input_kwargs={"X": add_5, "Y": add_2, "Z": add_2}
    )

    cube_2 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_3})

    store_named_attribute_1 = nw.new_node(
        Nodes.StoreNamedAttribute,
        input_kwargs={
            "Geometry": cube_2.outputs["Mesh"],
            "Name": "uv_map",
            3: cube_2.outputs["UV Map"],
        },
        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
    )

    instance_on_points = nw.new_node(
        Nodes.InstanceOnPoints,
        input_kwargs={"Points": transform_1, "Instance": store_named_attribute_1},
    )

    subtract_2 = nw.new_node(
        Nodes.Math, input_kwargs={0: add_4, 1: add}, attrs={"operation": "SUBTRACT"}
    )

    divide_2 = nw.new_node(
        Nodes.Math,
        input_kwargs={0: subtract_2, 1: add_3},
        attrs={"operation": "DIVIDE"},
    )

    grid_1 = nw.new_node(
        Nodes.MeshGrid,
        input_kwargs={
            "Size X": subtract_2,
            "Size Y": subtract,
            "Vertices X": divide_2,
            "Vertices Y": divide,
        },
    )

    store_named_attribute_2 = nw.new_node(
        Nodes.StoreNamedAttribute,
        input_kwargs={
            "Geometry": grid_1.outputs["Mesh"],
            "Name": "uv_map",
            3: grid_1.outputs["UV Map"],
        },
        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
    )

    transform_2 = nw.new_node(
        Nodes.Transform,
        input_kwargs={
            "Geometry": store_named_attribute_2,
            "Rotation": (1.5708, 0.0000, 0.0000),
        },
    )

    add_6 = nw.new_node(Nodes.Math, input_kwargs={0: add_1, 1: 0.1}, attrs={"operation": "ADD"})

    combine_xyz_4 = nw.new_node(
        Nodes.CombineXYZ, input_kwargs={"X": add_2, "Y": add_6, "Z": add_2}
    )

    cube_3 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_4})

    store_named_attribute_3 = nw.new_node(
        Nodes.StoreNamedAttribute,
        input_kwargs={
            "Geometry": cube_3.outputs["Mesh"],
            "Name": "uv_map",
            3: cube_3.outputs["UV Map"],
        },
        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
    )

    instance_on_points_1 = nw.new_node(
        Nodes.InstanceOnPoints,
        input_kwargs={"Points": transform_2, "Instance": store_named_attribute_3},
    )

    group_output = nw.new_node(
        Nodes.GroupOutput,
        input_kwargs={
            "Instances1": instance_on_points,
            "Instances2": instance_on_points_1,
        },
        attrs={"is_active_output": True},
    )


@node_utils.to_nodegroup(
    "nodegroup_handle_hole", singleton=False, type="GeometryNodeTree"
)
def nodegroup_handle_hole(nw: NodeWrangler):
    # Code generated using version 2.6.4 of the node_transpiler

    group_input = nw.new_node(
        Nodes.GroupInput,
        expose_input=[
            ("NodeSocketFloat", "X", 0.0000),
            ("NodeSocketFloat", "Z", 0.0000),
            ("NodeSocketFloat", "height", 0.5000),
            ("NodeSocketFloat", "hole_dist", 0.5000),
            ("NodeSocketInt", "Level", 0),
        ],
    )

    combine_xyz_3 = nw.new_node(
        Nodes.CombineXYZ,
        input_kwargs={
            "X": group_input.outputs["X"],
            "Y": 1.0000,
            "Z": group_input.outputs["Z"],
        },
    )

    cube_2 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_3})

    store_named_attribute = nw.new_node(
        Nodes.StoreNamedAttribute,
        input_kwargs={
            "Geometry": cube_2.outputs["Mesh"],
            "Name": "uv_map",
            3: cube_2.outputs["UV Map"],
        },
        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
    )

    subdivide_mesh_2 = nw.new_node(
        Nodes.SubdivideMesh, input_kwargs={"Mesh": store_named_attribute}
    )

    subdivision_surface_2 = nw.new_node(
        Nodes.SubdivisionSurface,
        input_kwargs={"Mesh": subdivide_mesh_2, "Level": group_input.outputs["Level"]},
    )

    multiply = nw.new_node(
        Nodes.Math,
        input_kwargs={0: group_input.outputs["height"]},
        attrs={"operation": "MULTIPLY"},
    )

    subtract = nw.new_node(
        Nodes.Math,
        input_kwargs={0: multiply, 1: group_input.outputs["hole_dist"]},
        attrs={"operation": "SUBTRACT"},
    )

    combine_xyz_4 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": subtract})

    transform_1 = nw.new_node(
        Nodes.Transform,
        input_kwargs={"Geometry": subdivision_surface_2, "Translation": combine_xyz_4},
    )

    group_output = nw.new_node(
        Nodes.GroupOutput,
        input_kwargs={"Geometry": transform_1},
        attrs={"is_active_output": True},
    )


def geometry_nodes(nw: NodeWrangler, **kwargs):
    # Code generated using version 2.6.4 of the node_transpiler

    depth = nw.new_node(Nodes.Value, label="depth")
    depth.outputs[0].default_value = kwargs["depth"]

    width = nw.new_node(Nodes.Value, label="width")
    width.outputs[0].default_value = kwargs["width"]

    height = nw.new_node(Nodes.Value, label="height")
    height.outputs[0].default_value = kwargs["height"]

    combine_xyz = nw.new_node(
        Nodes.CombineXYZ, input_kwargs={"X": depth, "Y": width, "Z": height}
    )

    cube = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz})

    store_named_attribute = nw.new_node(
        Nodes.StoreNamedAttribute,
        input_kwargs={
            "Geometry": cube.outputs["Mesh"],
            "Name": "uv_map",
            3: cube.outputs["UV Map"],
        },
        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
    )

    subdivide_mesh = nw.new_node(
        Nodes.SubdivideMesh, input_kwargs={"Mesh": store_named_attribute, "Level": 2}
    )

    sub_level = nw.new_node(Nodes.Integer, label="sub_level")
    sub_level.integer = kwargs["frame_sub_level"]

    subdivision_surface = nw.new_node(
        Nodes.SubdivisionSurface,
        input_kwargs={"Mesh": subdivide_mesh, "Level": sub_level},
    )

    differences = []

    if kwargs["has_handle"]:
        hole_depth = nw.new_node(Nodes.Value, label="hole_depth")
        hole_depth.outputs[0].default_value = kwargs["handle_depth"]

        hole_height = nw.new_node(Nodes.Value, label="hole_height")
        hole_height.outputs[0].default_value = kwargs["handle_height"]

        hole_dist = nw.new_node(Nodes.Value, label="hole_dist")
        hole_dist.outputs[0].default_value = kwargs["handle_dist_to_top"]

        handle_level = nw.new_node(Nodes.Integer, label="handle_level")
        handle_level.integer = kwargs["handle_sub_level"]
        handle_hole = nw.new_node(
            nodegroup_handle_hole().name,
            input_kwargs={
                "X": hole_depth,
                "Z": hole_height,
                "height": height,
                "hole_dist": hole_dist,
                "Level": handle_level,
            },
        )
        differences.append(handle_hole)

    thickness = nw.new_node(Nodes.Value, label="thickness")
    thickness.outputs[0].default_value = kwargs["thickness"]

    subtract = nw.new_node(
        Nodes.Math,
        input_kwargs={0: depth, 1: thickness},
        attrs={"operation": "SUBTRACT"},
    )

    subtract_1 = nw.new_node(
        Nodes.Math,
        input_kwargs={0: width, 1: thickness},
        attrs={"operation": "SUBTRACT"},
    )

    combine_xyz_1 = nw.new_node(
        Nodes.CombineXYZ, input_kwargs={"X": subtract, "Y": subtract_1, "Z": height}
    )

    cube_1 = nw.new_node(Nodes.MeshCube, input_kwargs={"Size": combine_xyz_1})

    store_named_attribute_1 = nw.new_node(
        Nodes.StoreNamedAttribute,
        input_kwargs={
            "Geometry": cube_1.outputs["Mesh"],
            "Name": "uv_map",
            3: cube_1.outputs["UV Map"],
        },
        attrs={"domain": "CORNER", "data_type": "FLOAT_VECTOR"},
    )

    subdivide_mesh_1 = nw.new_node(
        Nodes.SubdivideMesh, input_kwargs={"Mesh": store_named_attribute_1, "Level": 2}
    )

    subdivision_surface_1 = nw.new_node(
        Nodes.SubdivisionSurface,
        input_kwargs={"Mesh": subdivide_mesh_1, "Level": sub_level},
    )

    multiply = nw.new_node(
        Nodes.Math,
        input_kwargs={0: thickness, 2: 0.2500},
        attrs={"operation": "MULTIPLY"},
    )

    combine_xyz_2 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply})

    transform = nw.new_node(
        Nodes.Transform,
        input_kwargs={"Geometry": subdivision_surface_1, "Translation": combine_xyz_2},
    )

    if kwargs["has_holes"]:
        gap_size = nw.new_node(Nodes.Value, label="gap_size")
        gap_size.outputs[0].default_value = kwargs["hole_gap_size"]

        hole_edge_gap = nw.new_node(Nodes.Value, label="hole_edge_gap")
        hole_edge_gap.outputs[0].default_value = kwargs["hole_edge_gap"]

        hole_size = nw.new_node(Nodes.Value, label="hole_size")
        hole_size.outputs[0].default_value = kwargs["hole_size"]
        holes = nw.new_node(
            nodegroup_holes().name,
            input_kwargs={
                "height": height,
                "gap_size": gap_size,
                "hole_edge_gap": hole_edge_gap,
                "hole_size": hole_size,
                "depth": depth,
                "width": width,
            },
        )
        differences.extend([holes.outputs["Instances1"], holes.outputs["Instances2"]])

    difference = nw.new_node(
        Nodes.MeshBoolean,
        input_kwargs={
            "Mesh 1": subdivision_surface,
            "Mesh 2": [transform] + differences,
        },
    )

    realize_instances = nw.new_node(
        Nodes.RealizeInstances, input_kwargs={"Geometry": difference.outputs["Mesh"]}
    )

    multiply_1 = nw.new_node(
        Nodes.Math, input_kwargs={0: height}, attrs={"operation": "MULTIPLY"}
    )

    combine_xyz_3 = nw.new_node(Nodes.CombineXYZ, input_kwargs={"Z": multiply_1})

    transform_geometry = nw.new_node(
        Nodes.Transform,
        input_kwargs={"Geometry": realize_instances, "Translation": combine_xyz_3},
    )

    set_material = nw.new_node(
        Nodes.SetMaterial,
        input_kwargs={
            "Geometry": transform_geometry,
            "Material": surface.shaderfunc_to_material(shader_rough_plastic),
        },
    )

    group_output = nw.new_node(
        Nodes.GroupOutput,
        input_kwargs={"Geometry": set_material},
        attrs={"is_active_output": True},
    )


class BasketBaseFactory(AssetFactory):
    def __init__(self, factory_seed, coarse=False):
        super(BasketBaseFactory, self).__init__(factory_seed, coarse=coarse)
        self.params = self.get_asset_params()
        self.seed = factory_seed
        self.get_params_dict()
    
    def get_params_dict(self):
        self.params_dict = {
            "depth": ['continuous', (0.1, 0.6)],
            "width": ['continuous', (0.1, 0.7)],
            "height": ['continuous', (0.05, 0.4)],
            "frame_sub_level": ['discrete', [0, 3]],
            "thickness": ['continuous', (0.001, 0.03)],
            "has_handle": ['discrete', [0, 1]],
            "handle_sub_level": ['discrete', [0, 1, 2]],
            "handle_depth": ['continuous', (0.2, 0.6)],
            "handle_height": ['continuous', (0.1, 0.3)],
            "handle_dist_to_top": ['continuous', (0.08, 0.4)],
            "has_holes": ['discrete', [0, 1]],
            "hole_gap_size": ['continuous', (0.5, 2.0)],
            "hole_edge_gap": ['continuous', (0.04, 0.1)],
            "hole_size": ['continuous', (0.007, 0.02)]
        }

    def fix_unused_params(self, params):
        if params['height'] < 0.12:
            params['has_holes'] = 0
        if params['has_handle'] == 0:
            params["handle_sub_level"] = 1
            params["handle_depth"] = 0.3
            params["handle_height"] = 0.2
            params["handle_dist_to_top"] = 0.115
        if params['has_holes'] == 0:
            params["hole_gap_size"] = 0.95
            params["hole_edge_gap"] = 0.05
            params["hole_size"] = 0.0075
        return params

    def update_params(self, params):
        # TODO: to allow random material
        self.seed = int(1000 * time.time()) % 2**32
        
        handle_depth = params['depth'] * params['handle_depth']
        handle_height = params['height'] * params['handle_height']
        handle_dist_to_top = handle_height * 0.5 + params['height'] * params["handle_dist_to_top"]
        if params['height'] < 0.12:
            params["has_holes"] = 0
        hole_gap_size = params['hole_size'] * params["hole_gap_size"]
        parameters = {
            "depth": params["depth"],
            "width": params["width"],
            "height": params["height"],
            "frame_sub_level": params["frame_sub_level"],
            "thickness": params["thickness"],
            "has_handle": params["has_handle"] > 0,
            "handle_sub_level": params["handle_sub_level"],
            "handle_depth": handle_depth,
            "handle_height": handle_height,
            "handle_dist_to_top": handle_dist_to_top,
            "has_holes": params["has_holes"] > 0,
            "hole_gap_size": hole_gap_size,
            "hole_edge_gap": params["hole_edge_gap"],
            "hole_size": params["hole_size"],
        }
        self.params.update(parameters)

    def get_asset_params(self, i=0):
        params = {}
        if params.get("depth", None) is None:
            params["depth"] = uniform(0.15, 0.4)
        if params.get("width", None) is None:
            params["width"] = uniform(0.2, 0.6)
        if params.get("height", None) is None:
            params["height"] = uniform(0.06, 0.24)
        if params.get("frame_sub_level", None) is None:
            params["frame_sub_level"] = np.random.choice([0, 3], p=[0.5, 0.5])
        if params.get("thickness", None) is None:
            params["thickness"] = uniform(0.001, 0.005)

        if params.get("has_handle", None) is None:
            params["has_handle"] = np.random.choice([True, False], p=[0.8, 0.2])
        if params.get("handle_sub_level", None) is None:
            params["handle_sub_level"] = np.random.choice([0, 1, 2], p=[0.2, 0.4, 0.4])
        if params.get("handle_depth", None) is None:
            params["handle_depth"] = params["depth"] * uniform(0.2, 0.4)
        if params.get("handle_height", None) is None:
            params["handle_height"] = params["height"] * uniform(0.1, 0.25)
        if params.get("handle_dist_to_top", None) is None:
            params["handle_dist_to_top"] = params["handle_height"] * 0.5 + params[
                "height"
            ] * uniform(0.08, 0.15)

        if params.get("has_holes", None) is None:
            if params["height"] < 0.12:
                params["has_holes"] = False
            else:
                params["has_holes"] = np.random.choice([True, False], p=[0.5, 0.5])
        if params.get("hole_size", None) is None:
            params["hole_size"] = uniform(0.005, 0.01)
        if params.get("hole_gap_size", None) is None:
            params["hole_gap_size"] = params["hole_size"] * uniform(0.8, 1.1)
        if params.get("hole_edge_gap", None) is None:
            params["hole_edge_gap"] = uniform(0.04, 0.06)

        return params

    def create_asset(self, i=0, **params):
        bpy.ops.mesh.primitive_plane_add(
            size=1,
            enter_editmode=False,
            align="WORLD",
            location=(0, 0, 0),
            scale=(1, 1, 1),
        )
        obj = bpy.context.active_object
        np.random.seed(self.seed)
        random.seed(self.seed)
        
        surface.add_geomod(
            obj, geometry_nodes, attributes=[], apply=True, input_kwargs=self.params
        )
        tagging.tag_system.relabel_obj(obj)
        return obj