wabbajack/Wabbajack.Compiler/CompilationSteps/MatchSimilarTextures.cs

using System.Linq;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using Wabbajack.DTOs;
using Wabbajack.DTOs.Directives;
using Wabbajack.Hashing.PHash;
using Wabbajack.Paths;
using Wabbajack.VFS;

namespace Wabbajack.Compiler.CompilationSteps
{
    public class MatchSimilarTextures : ACompilationStep
    {
        private ILookup<RelativePath, VirtualFile> _byName;
        public MatchSimilarTextures(ACompiler compiler) : base(compiler)
        {
            _byName = _compiler.IndexedFiles.SelectMany(kv => kv.Value)
                .Where(f => f.Name.FileName.Extension == DDS)
                .Where(f => f.ImageState != default)
                .ToLookup(f => f.Name.FileName.FileNameWithoutExtension);
        }

        private const float PerceptualTolerance = 0.80f;

        private static Extension DDS = new(".dds");


        private static string[] PostFixes = new[] {"_n", "_d", "_s"};
        public override async ValueTask<Directive?> Run(RawSourceFile source)
        {
            if (source.File.Name.FileName.Extension == DDS && source.File.ImageState != null)
            {
                _compiler._logger.LogInformation("Looking for texture match for {source}", source.File.FullPath);
                (float Similarity, VirtualFile File) found = _byName[source.Path.FileNameWithoutExtension]
                    .Select(f => (ImageLoader.ComputeDifference(f.ImageState!.PerceptualHash, source.File.ImageState.PerceptualHash), f))
                    .Select(f =>
                    {
                        return f;
                    })
                    .OrderByDescending(f => f.Item1)
                    .FirstOrDefault();

                if (found == default || found.Similarity <= PerceptualTolerance)
                {
                    // This looks bad, but it's fairly simple: normal and displacement textures don't match very well
                    // via perceptual hashing. So instead we'll try to find a diffuse map with the same name, and look
                    // for normal maps in the same folders. Example: roof_n.dds didn't match, so find a match betweeen
                    // roof.dds and a perceptual match in the downloads. Then try to find a roof_n.dds in the same folder
                    // as the match we found for roof.dds. 
                    found = default;
                    var r = from postfix in PostFixes
                            where source.File.Name.FileName.FileNameWithoutExtension.EndsWith(postfix)
                            let mainFileName =
                                source.File.Name.FileName.FileNameWithoutExtension.ToString()[..^postfix.Length] +
                                ".dds"
                            let mainFile = source.File.InSameFolder(mainFileName.ToRelativePath())
                        where mainFile != null
                        from mainMatch in _byName[mainFile.FullPath.FileName.FileNameWithoutExtension]
                        where mainMatch.ImageState != null
                        where mainFile.ImageState != null
                            let similarity = ImageLoader.ComputeDifference(mainFile.ImageState.PerceptualHash, mainMatch.ImageState.PerceptualHash)
                        where similarity >= PerceptualTolerance
                        orderby similarity descending
                        let foundFile = mainMatch.InSameFolder(source.Path.FileName)
                            where foundFile != null
                        select (similarity, postfix, mainFile, mainMatch, foundFile);

                    foreach (var record in r)
                    {
                        _compiler._logger.LogInformation("Found Match for {source} {data}", source.File.Name, record.foundFile.ImageState);
                        
                    }

                    var foundRec = r.FirstOrDefault();
                    if (foundRec == default)
                    {
                        return null;
                    }

                    found = (foundRec.similarity, foundRec.foundFile);
                }
                
                _compiler._logger.LogInformation("Found Match for {source} {sourceData} {destData}", source.File.Name, source.File.ImageState, found.File.ImageState);

                var rv = source.EvolveTo<TransformedTexture>();
                rv.ArchiveHashPath = found.File.MakeRelativePaths();
                rv.ImageState = source.File.ImageState!;
                
                return rv;
            }

            return null;
        }
    }
}