blender
/
blender-addons
17
34
Fork 94
Code Issues 529 Pull Requests 32 Projects 1 Releases Wiki Activity

Node Wrangler: improve Add Principled Setup file matching logic 

Improve how "Add Principled Setup" decides which selected files should go into
which sockets.

The new tests verify that the new logic does the right thing with metallic
textures downloaded from here, plus some corner cases. None of these worked
before:
* https://poliigon.com
* https://ambientcg.com
* https://3dtextures.me
* https://polyhaven.com

Adding new tests for textures from other sites should be simple by just
extending the existing tests.

What the new code does:
* Converts file names into tag lists using `split_into_components()`
* Remove common prefixes from all tag lists
* Remove common suffixes from all tag lists
* Ignore files matching no socket tags (filters out README files and
  previews usually)
* Iterate ^ until nothing changes any more
* Do the same matching as before in `match_files_to_socket_names()`,
  but on the filtered tag lists

Other changes:
* `node_wrangler.py` was moved into a `node_wrangler/main.py` to enable putting
  tests next to it. Inspired by `io_curve_svg/` and its `svg_util_test.py`.
* File-names-to-socket-matching code was moved into its own file `util.py`
  so that both tests and the production code can find it
* Tests were added in `node_wrangler/util_test.py`

Differential Revision: https://developer.blender.org/D16940
This commit is contained in:
Johan Walles 2023-01-19 16:58:39 +01:00 committed by Brecht Van Lommel
parent 6fcd157f24
commit c0a678d368
3 changed files with 399 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
node_wrangler.py → node_wrangler/__init__.py
View File

@ -27,6 +27,7 @@ from bpy.props import (
from bpy_extras.io_utils import ImportHelper, ExportHelper
from gpu_extras.batch import batch_for_shader
from mathutils import Vector
from .util import match_files_to_socket_names, split_into_components
from nodeitems_utils import node_categories_iter, NodeItemCustom
from math import cos, sin, pi, hypot
from os import path
@ -679,7 +680,7 @@ def get_output_location(tree):
class NWPrincipledPreferences(bpy.types.PropertyGroup):
base_color: StringProperty(
name='Base Color',
default='diffuse diff albedo base col color',
default='diffuse diff albedo base col color basecolor',
description='Naming Components for Base Color maps')
sss_color: StringProperty(
name='Subsurface Color',
@ -2712,25 +2713,6 @@ class NWAddPrincipledSetup(Operator, NWBase, ImportHelper):
self.report({'INFO'}, 'Select Principled BSDF')
return {'CANCELLED'}
# Helper_functions
def split_into__components(fname):
# Split filename into components
# 'WallTexture_diff_2k.002.jpg' -> ['Wall', 'Texture', 'diff', 'k']
# Remove extension
fname = path.splitext(fname)[0]
# Remove digits
fname = ''.join(i for i in fname if not i.isdigit())
# Separate CamelCase by space
fname = re.sub(r"([a-z])([A-Z])", r"\g<1> \g<2>",fname)
# Replace common separators with SPACE
separators = ['_', '.', '-', '__', '--', '#']
for sep in separators:
fname = fname.replace(sep, ' ')
components = fname.split(' ')
components = [c.lower() for c in components]
return components
# Filter textures names for texturetypes in filenames
# [Socket Name, [abbreviations and keyword list], Filename placeholder]
tags = context.preferences.addons[__name__].preferences.principled_tags
@ -2752,19 +2734,7 @@ class NWAddPrincipledSetup(Operator, NWBase, ImportHelper):
['Ambient Occlusion', tags.ambient_occlusion.split(' '), None],
]
# Look through texture_types and set value as filename of first matched file
def match_files_to_socket_names():
for sname in socketnames:
for file in self.files:
fname = file.name
filenamecomponents = split_into__components(fname)
matches = set(sname[1]).intersection(set(filenamecomponents))
# TODO: ignore basename (if texture is named "fancy_metal_nor", it will be detected as metallic map, not normal map)
if matches:
sname[2] = fname
break
match_files_to_socket_names()
match_files_to_socket_names(self.files, socketnames)
# Remove socketnames without found files
socketnames = [s for s in socketnames if s[2]
and path.exists(self.directory+s[2])]
@ -2838,7 +2808,7 @@ class NWAddPrincipledSetup(Operator, NWBase, ImportHelper):
# NORMAL NODES
if sname[0] == 'Normal':
# Test if new texture node is normal or bump map
fname_components = split_into__components(sname[2])
fname_components = split_into_components(sname[2])
match_normal = set(normal_abbr).intersection(set(fname_components))
match_bump = set(bump_abbr).intersection(set(fname_components))
if match_normal:
@ -2855,7 +2825,7 @@ class NWAddPrincipledSetup(Operator, NWBase, ImportHelper):
elif sname[0] == 'Roughness':
# Test if glossy or roughness map
fname_components = split_into__components(sname[2])
fname_components = split_into_components(sname[2])
match_rough = set(rough_abbr).intersection(set(fname_components))
match_gloss = set(gloss_abbr).intersection(set(fname_components))

164
node_wrangler/util.py Normal file
View File

@ -0,0 +1,164 @@
# SPDX-License-Identifier: GPL-2.0-or-later
from os import path
import re
def split_into_components(fname):
"""
Split filename into components
'WallTexture_diff_2k.002.jpg' -> ['Wall', 'Texture', 'diff', 'k']
"""
# Remove extension
fname = path.splitext(fname)[0]
# Remove digits
fname = "".join(i for i in fname if not i.isdigit())
# Separate CamelCase by space
fname = re.sub(r"([a-z])([A-Z])", r"\g<1> \g<2>", fname)
# Replace common separators with SPACE
separators = ["_", ".", "-", "__", "--", "#"]
for sep in separators:
fname = fname.replace(sep, " ")
components = fname.split(" ")
components = [c.lower() for c in components]
return components
def remove_common_prefix(names_to_tag_lists):
"""
Accepts a mapping of file names to tag lists that should be used for socket
matching.
This function modifies the provided mapping so that any common prefix
between all the tag lists is removed.
Returns true if some prefix was removed, false otherwise.
"""
if not names_to_tag_lists:
return False
sample_tags = next(iter(names_to_tag_lists.values()))
if not sample_tags:
return False
common_prefix = sample_tags[0]
for tag_list in names_to_tag_lists.values():
if tag_list[0] != common_prefix:
return False
for name, tag_list in names_to_tag_lists.items():
names_to_tag_lists[name] = tag_list[1:]
return True
def remove_common_suffix(names_to_tag_lists):
"""
Accepts a mapping of file names to tag lists that should be used for socket
matching.
This function modifies the provided mapping so that any common suffix
between all the tag lists is removed.
Returns true if some suffix was removed, false otherwise.
"""
if not names_to_tag_lists:
return False
sample_tags = next(iter(names_to_tag_lists.values()))
if not sample_tags:
return False
common_suffix = sample_tags[-1]
for tag_list in names_to_tag_lists.values():
if tag_list[-1] != common_suffix:
return False
for name, tag_list in names_to_tag_lists.items():
names_to_tag_lists[name] = tag_list[:-1]
return True
def files_to_clean_file_names_for_sockets(files, sockets):
"""
Accepts a list of files and a list of sockets.
Returns a mapping from file names to tag lists that should be used for
classification.
A file is something that we can do x.name on to figure out the file name.
A socket is a tuple containing:
* name
* list of tags
* a None field where the selected file name will go later. Ignored by us.
"""
names_to_tag_lists = {}
for file in files:
names_to_tag_lists[file.name] = split_into_components(file.name)
all_tags = set()
for socket in sockets:
socket_tags = socket[1]
all_tags.update(socket_tags)
while True:
something_changed = False
# Common prefixes / suffixes provide zero information about what file
# should go to which socket, but they can confuse the mapping. So we get
# rid of them here.
something_changed |= remove_common_prefix(names_to_tag_lists)
something_changed |= remove_common_suffix(names_to_tag_lists)
# Names matching zero tags provide no value, remove those
names_to_remove = set()
for name, tag_list in names_to_tag_lists.items():
match_found = False
for tag in tag_list:
if tag in all_tags:
match_found = True
if not match_found:
names_to_remove.add(name)
for name_to_remove in names_to_remove:
del names_to_tag_lists[name_to_remove]
something_changed = True
if not something_changed:
break
return names_to_tag_lists
def match_files_to_socket_names(files, sockets):
"""
Given a list of files and a list of sockets, match file names to sockets.
A file is something that you can get a file name out of using x.name.
After this function returns, all possible sockets have had their file names
filled in. Sockets without any matches will not get their file names
changed.
Sockets list format. Note that all file names are initially expected to be
None. Tags are strings, as are the socket names: [
[
socket_name, [tags], Optional[file_name]
]
]
"""
names_to_tag_lists = files_to_clean_file_names_for_sockets(files, sockets)
for sname in sockets:
for name, tag_list in names_to_tag_lists.items():
if sname[0] == "Normal" and "dx" in tag_list:
# Blender wants GL normals, not DX (DirectX) ones:
# https://www.reddit.com/r/blender/comments/rbuaua/texture_contains_normaldx_and_normalgl_files/
continue
matches = set(sname[1]).intersection(set(tag_list))
if matches:
sname[2] = name
break

230
node_wrangler/util_test.py Executable file
View File

@ -0,0 +1,230 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later
# pylint: disable=missing-function-docstring
# pylint: disable=missing-class-docstring
import unittest
from dataclasses import dataclass
# XXX Not really nice, but that hack is needed to allow execution of that test
# from both automated CTest and by directly running the file manually.
if __name__ == "__main__":
from util import match_files_to_socket_names
else:
from .util import match_files_to_socket_names
# From NWPrincipledPreferences 2023-01-06
TAGS_DISPLACEMENT = "displacement displace disp dsp height heightmap".split(" ")
TAGS_BASE_COLOR = "diffuse diff albedo base col color basecolor".split(" ")
TAGS_SUBSURFACE_COLOR = "sss subsurface".split(" ")
TAGS_METALLIC = "metallic metalness metal mtl".split(" ")
TAGS_SPECULAR = "specularity specular spec spc".split(" ")
TAGS_ROUGHNESS = "roughness rough rgh".split(" ")
TAGS_GLOSS = "gloss glossy glossiness".split(" ")
TAGS_NORMAL = "normal nor nrm nrml norm".split(" ")
TAGS_BUMP = "bump bmp".split(" ")
TAGS_TRANSMISSION = "transmission transparency".split(" ")
TAGS_EMISSION = "emission emissive emit".split(" ")
TAGS_ALPHA = "alpha opacity".split(" ")
TAGS_AMBIENT_OCCLUSION = "ao ambient occlusion".split(" ")
@dataclass
class MockFile:
name: str
def sockets_fixture():
return [
["Displacement", TAGS_DISPLACEMENT, None],
["Base Color", TAGS_BASE_COLOR, None],
["Subsurface Color", TAGS_SUBSURFACE_COLOR, None],
["Metallic", TAGS_METALLIC, None],
["Specular", TAGS_SPECULAR, None],
["Roughness", TAGS_ROUGHNESS + TAGS_GLOSS, None],
["Normal", TAGS_NORMAL + TAGS_BUMP, None],
["Transmission", TAGS_TRANSMISSION, None],
["Emission", TAGS_EMISSION, None],
["Alpha", TAGS_ALPHA, None],
["Ambient Occlusion", TAGS_AMBIENT_OCCLUSION, None],
]
def assert_sockets(asserter, sockets, expected):
checked_sockets = set()
errors = []
for socket_name, expected_path in expected.items():
if isinstance(expected_path, str):
expected_path = [expected_path]
socket_found = False
for socket in sockets:
if socket[0] != socket_name:
continue
socket_found = True
actual_path = socket[2]
if actual_path not in expected_path:
errors.append(
f"{socket_name:12}: Got {actual_path} but expected {expected_path}"
)
checked_sockets.add(socket_name)
break
asserter.assertTrue(socket_found)
asserter.assertCountEqual([], errors)
for socket in sockets:
if socket[0] in checked_sockets:
continue
asserter.assertEqual(socket[2], None)
class TestPutFileNamesInSockets(unittest.TestCase):
def test_no_files_selected(self):
sockets = sockets_fixture()
match_files_to_socket_names([], sockets)
assert_sockets(self, sockets, {})
def test_weird_filename(self):
sockets = sockets_fixture()
match_files_to_socket_names(
[MockFile(""), MockFile(".jpg"), MockFile(" .png"), MockFile("...")],
sockets,
)
assert_sockets(self, sockets, {})
def test_poliigon(self):
"""Texture from: https://www.poliigon.com/texture/metal-spotty-discoloration-001/3225"""
# NOTE: These files all have directory prefixes. That's on purpose. Files
# without directory prefixes are tested in test_ambientcg_metal().
files = [
MockFile("d/MetalSpottyDiscoloration001_COL_2K_METALNESS.jpg"),
MockFile("d/MetalSpottyDiscoloration001_Cube.jpg"),
MockFile("d/MetalSpottyDiscoloration001_DISP16_2K_METALNESS.tif"),
MockFile("d/MetalSpottyDiscoloration001_DISP_2K_METALNESS.jpg"),
MockFile("d/MetalSpottyDiscoloration001_Flat.jpg"),
MockFile("d/MetalSpottyDiscoloration001_METALNESS_2K_METALNESS.jpg"),
MockFile("d/MetalSpottyDiscoloration001_NRM16_2K_METALNESS.tif"),
MockFile("d/MetalSpottyDiscoloration001_NRM_2K_METALNESS.jpg"),
MockFile("d/MetalSpottyDiscoloration001_ROUGHNESS_2K_METALNESS.jpg"),
MockFile("d/MetalSpottyDiscoloration001_Sphere.jpg"),
]
sockets = sockets_fixture()
match_files_to_socket_names(files, sockets)
assert_sockets(
self,
sockets,
{
"Base Color": "d/MetalSpottyDiscoloration001_COL_2K_METALNESS.jpg",
"Displacement": [
"d/MetalSpottyDiscoloration001_DISP16_2K_METALNESS.tif",
"d/MetalSpottyDiscoloration001_DISP_2K_METALNESS.jpg",
],
"Metallic": "d/MetalSpottyDiscoloration001_METALNESS_2K_METALNESS.jpg",
"Normal": [
"d/MetalSpottyDiscoloration001_NRM16_2K_METALNESS.tif",
"d/MetalSpottyDiscoloration001_NRM_2K_METALNESS.jpg",
],
"Roughness": "d/MetalSpottyDiscoloration001_ROUGHNESS_2K_METALNESS.jpg",
},
)
def test_ambientcg(self):
"""Texture from: https://ambientcg.com/view?id=MetalPlates003"""
# NOTE: These files have no directory prefix. That's on purpose. Files
# with directory prefixes are tested in test_poliigon_metal().
files = [
MockFile("MetalPlates001_1K-JPG.usda"),
MockFile("MetalPlates001_1K-JPG.usdc"),
MockFile("MetalPlates001_1K_Color.jpg"),
MockFile("MetalPlates001_1K_Displacement.jpg"),
MockFile("MetalPlates001_1K_Metalness.jpg"),
MockFile("MetalPlates001_1K_NormalDX.jpg"),
MockFile("MetalPlates001_1K_NormalGL.jpg"),
MockFile("MetalPlates001_1K_Roughness.jpg"),
MockFile("MetalPlates001_PREVIEW.jpg"),
]
sockets = sockets_fixture()
match_files_to_socket_names(files, sockets)
assert_sockets(
self,
sockets,
{
"Base Color": "MetalPlates001_1K_Color.jpg",
"Displacement": "MetalPlates001_1K_Displacement.jpg",
"Metallic": "MetalPlates001_1K_Metalness.jpg",
# Blender wants GL normals:
# https://www.reddit.com/r/blender/comments/rbuaua/texture_contains_normaldx_and_normalgl_files/
"Normal": "MetalPlates001_1K_NormalGL.jpg",
"Roughness": "MetalPlates001_1K_Roughness.jpg",
},
)
def test_3dtextures_me(self):
"""Texture from: https://3dtextures.me/2022/05/13/metal-006/"""
files = [
MockFile("Material_2079.jpg"),
MockFile("Metal_006_ambientOcclusion.jpg"),
MockFile("Metal_006_basecolor.jpg"),
MockFile("Metal_006_height.png"),
MockFile("Metal_006_metallic.jpg"),
MockFile("Metal_006_normal.jpg"),
MockFile("Metal_006_roughness.jpg"),
]
sockets = sockets_fixture()
match_files_to_socket_names(files, sockets)
assert_sockets(
self,
sockets,
{
"Ambient Occlusion": "Metal_006_ambientOcclusion.jpg",
"Base Color": "Metal_006_basecolor.jpg",
"Displacement": "Metal_006_height.png",
"Metallic": "Metal_006_metallic.jpg",
"Normal": "Metal_006_normal.jpg",
"Roughness": "Metal_006_roughness.jpg",
},
)
def test_polyhaven(self):
"""Texture from: https://polyhaven.com/a/rusty_metal_02"""
files = [
MockFile("rusty_metal_02_ao_1k.jpg"),
MockFile("rusty_metal_02_arm_1k.jpg"),
MockFile("rusty_metal_02_diff_1k.jpg"),
MockFile("rusty_metal_02_disp_1k.png"),
MockFile("rusty_metal_02_nor_dx_1k.exr"),
MockFile("rusty_metal_02_nor_gl_1k.exr"),
MockFile("rusty_metal_02_rough_1k.exr"),
MockFile("rusty_metal_02_spec_1k.png"),
]
sockets = sockets_fixture()
match_files_to_socket_names(files, sockets)
assert_sockets(
self,
sockets,
{
"Ambient Occlusion": "rusty_metal_02_ao_1k.jpg",
"Base Color": "rusty_metal_02_diff_1k.jpg",
"Displacement": "rusty_metal_02_disp_1k.png",
"Normal": "rusty_metal_02_nor_gl_1k.exr",
"Roughness": "rusty_metal_02_rough_1k.exr",
"Specular": "rusty_metal_02_spec_1k.png",
},
)
if __name__ == "__main__":
unittest.main(verbosity=2)