make uv detection just a tad bit more correct, fix uvgen

Lilaa3 · Oct 26, 2024 · 03feb19 · 03feb19
1 parent 7dcef32
commit 03feb19
Showing 1 changed file with 43 additions and 37 deletions.
diff --git a/fast64_internal/f3d/bsdf_converter/converter.py b/fast64_internal/f3d/bsdf_converter/converter.py
@@ -580,13 +580,6 @@ def bsdf_mat_to_abstracted(material: Material):
         color_inp = next(("Color" for inp in color_shader.inputs if inp.name == "Color"), None)
     assert color_inp is not None, f"Could not find color input in {material.name}"
 
-    # uv map
-    found_uv_map_nodes = find_linked_nodes(output_node, lambda node: node.bl_idname == "ShaderNodeUVMap")
-    found_uv_map_names = list(dict.fromkeys([node.uv_map for node in found_uv_map_nodes]).keys())
-    if len(found_uv_map_names) > 1:
-        print(f"WARNING: More than 1 UV map being used in {material.name}. Using first UV map.")
-    abstracted_mat.uv_map = found_uv_map_names[0] if len(found_uv_map_names) > 0 else ""
-
     # vertex colors
     def get_vtx_layer(nodes):
         layer_names = list(dict.fromkeys([node.layer_name for node in nodes]).keys())
@@ -621,6 +614,7 @@ def get_vtx_layer(nodes):
     abstracted_mat.alpha_is_median = len(mean_vtx_alpha_nodes) > 0
 
     # textures
+    found_uv_map_nodes = []
     alpha_textures = find_linked_nodes(
         alpha_shader,
         lambda node: node.bl_idname == "ShaderNodeTexImage",
@@ -634,34 +628,46 @@ def get_vtx_layer(nodes):
     textures: list[ShaderNodeTexImage] = list(dict.fromkeys(color_textures + alpha_textures).keys())
     if len(textures) > 2:
         print(f"WARNING: More than 2 textures connected to {material.name}.")
-    if len(textures) > 0:
-        for tex_node in textures[:2]:
-            abstracted_tex = AbstractedN64Texture(tex_node.image)
-            mapping = find_linked_nodes(tex_node, lambda node: node.bl_idname == "ShaderNodeMapping")
-            if len(mapping) > 1:
-                print(f"WARNING: More than 1 mapping node connected to {tex_node.name}.")
-            elif len(mapping) == 1:
-                mapping = mapping[0]
-                abstracted_tex.offset = tuple(mapping.inputs["Location"].default_value)
-                abstracted_tex.scale = tuple(mapping.inputs["Scale"].default_value)
-            uv_gen = find_linked_nodes(
+    for tex_node in textures[:2]:
+        abstracted_tex = AbstractedN64Texture(tex_node.image)
+        found_uv_map_nodes.extend(
+            find_linked_nodes(
                 tex_node,
-                lambda node: node.bl_idname == "ShaderNodeTexCoord",
+                lambda node: node.bl_idname == "ShaderNodeUVMap",
                 specific_input_sockets={"Vector"},
-                specific_output_sockets={"Normal"},
+                specific_output_sockets={"UV"},
             )
-            if uv_gen:
-                abstracted_mat.uv_gen = True
-            if tex_node.interpolation == "Closest":
-                abstracted_mat.point_filtering = True
-            abstracted_tex.repeat = tex_node.extension == "REPEAT"
-            abstracted_tex.set_color = tex_node in color_textures
-            abstracted_tex.set_alpha = tex_node in alpha_textures
-            if abstracted_tex.set_color:
-                abstracted_mat.texture_sets_col = True
-            if abstracted_tex.set_alpha:
-                abstracted_mat.texture_sets_alpha = True
-            abstracted_mat.textures.append(abstracted_tex)
+        )
+        mapping = find_linked_nodes(tex_node, lambda node: node.bl_idname == "ShaderNodeMapping")
+        if len(mapping) > 1:
+            print(f"WARNING: More than 1 mapping node connected to {tex_node.name}.")
+        elif len(mapping) == 1:
+            mapping = mapping[0]
+            abstracted_tex.offset = tuple(mapping.inputs["Location"].default_value)
+            abstracted_tex.scale = tuple(mapping.inputs["Scale"].default_value)
+        uv_gen = find_linked_nodes(
+            tex_node,
+            lambda node: node.bl_idname == "ShaderNodeTexCoord",
+            specific_input_sockets={"Vector"},
+            specific_output_sockets={"Camera", "Window", "Reflection"},
+        )
+        if uv_gen:
+            abstracted_mat.uv_gen = True
+        if tex_node.interpolation == "Closest":
+            abstracted_mat.point_filtering = True
+        abstracted_tex.repeat = tex_node.extension == "REPEAT"
+        abstracted_tex.set_color = tex_node in color_textures
+        abstracted_tex.set_alpha = tex_node in alpha_textures
+        if abstracted_tex.set_color:
+            abstracted_mat.texture_sets_col = True
+        if abstracted_tex.set_alpha:
+            abstracted_mat.texture_sets_alpha = True
+        abstracted_mat.textures.append(abstracted_tex)
+    found_uv_map_names = list(dict.fromkeys([node.uv_map for node in found_uv_map_nodes]).keys())
+    if len(found_uv_map_names) > 1:
+        print(f"WARNING: More than 1 UV map being used in {material.name}. Using first UV map.")
+    abstracted_mat.uv_map = found_uv_map_names[0] if len(found_uv_map_names) > 0 else ""
+
     if not color_textures:
         abstracted_mat.color.r, abstracted_mat.color.g, abstracted_mat.color.b = color_shader.inputs[
             color_inp
@@ -814,7 +820,7 @@ def obj_to_f3d(
             return True
     print(f"Converting BSDF materials in {obj.name}")
     uvs = np.empty((len(obj.data.loops), 2), dtype=np.float32) if len(obj.data.uv_layers) != 1 else None
-    colors = np.ones((len(obj.data.loops), 4), dtype=np.float32) if len(obj.data.color_attributes) != 0 else None
+    colors = np.ones((len(obj.data.loops), 4), dtype=np.float32)
 
     loop_indexes: dict[Material, list] = {}
     for poly in obj.data.polygons:
@@ -827,12 +833,12 @@ def obj_to_f3d(
             loop_indexes[material].append(loop_idx)
 
     def get_layer_and_convert(layer_name: str | None):
-        layer = obj.data.color_attributes.get(layer_name or "", obj.data.color_attributes.active)
+        layer = obj.data.vertex_colors.get(layer_name or "", obj.data.vertex_colors.active)
         if layer is None:
             return None
         layer_name = layer.name
         convertColorAttribute(obj.data, layer_name)
-        return obj.data.color_attributes[layer_name]  # HACK: layer cannot be trusted
+        return obj.data.vertex_colors[layer_name]  # HACK: layer cannot be trusted
 
     for index, material_slot in enumerate(obj.material_slots):
         material = material_slot.material
@@ -873,8 +879,8 @@ def get_layer_and_convert(layer_name: str | None):
         obj.data.uv_layers.active = obj.data.uv_layers["UVMap"]
         obj.data.uv_layers["UVMap"].data.foreach_set("uv", uvs.flatten())
 
-    while len(obj.data.color_attributes) > 0:  # remove all existing colors
-        obj.data.color_attributes.remove(obj.data.color_attributes[0])
+    while len(obj.data.vertex_colors) > 0:  # remove all existing colors
+        obj.data.vertex_colors.remove(obj.data.vertex_colors[0])
     # get the alpha as rgb, then flatten it
     alpha_layer = obj.data.color_attributes.new("Alpha", "FLOAT_COLOR", "CORNER")
     alpha_layer.data.foreach_set("color", np.repeat(colors[:, 3][:, np.newaxis], 4, axis=1).flatten())