Merge pull request #1538 from wabbajack-tools/fix-texture-matching

Rework texture matching again, support for _n and _d files, and fixin…

Wabbajack.Common/Consts.cs

@@ -167,12 +167,12 @@ namespace Wabbajack.Common
         public static RelativePath ModListTxt = (RelativePath)"modlist.txt";
         public static RelativePath ModOrganizer2Exe = (RelativePath)"ModOrganizer.exe";
         public static RelativePath ModOrganizer2Ini = (RelativePath)"ModOrganizer.ini";
-        public static string AuthorAPIKeyFile = "author-api-key.txt";
+        public static readonly string AuthorAPIKeyFile = "author-api-key.txt";
 
-        public static Uri WabbajackOrg = new Uri("https://www.wabbajack.org/");
+        public static readonly Uri WabbajackOrg = new Uri("https://www.wabbajack.org/");
         
-        public static long UPLOADED_FILE_BLOCK_SIZE = (long)1024 * 1024 * 2;
+        public static readonly long UploadedFileBlockSize = (long)1024 * 1024 * 2;
 
-        public static HashSet<Extension> TextureExtensions = new() {new Extension(".dds"), new Extension(".tga")};
+        public static readonly HashSet<Extension> TextureExtensions = new() {new Extension(".dds"), new Extension(".tga")};
     }
 }

Wabbajack.Common/Paths/FullPath.cs

@@ -72,5 +72,25 @@ namespace Wabbajack.Common
         }
 
         public RelativePath FileName => Paths.Length == 0 ? Base.FileName : Paths.Last().FileName;
+
+        /// <summary>
+        /// Creates a new full path, with relativePath combined with the deepest leaf in the full path
+        /// </summary>
+        /// <param name="relativePath"></param>
+        /// <returns></returns>
+        public FullPath InSameFolder(RelativePath relativePath)
+        {
+            if (Paths.Length == 0)
+            {
+                return new FullPath(Base.Parent.Combine(relativePath));
+            }
+            else
+            {
+                var paths = new RelativePath[Paths.Length];
+                Paths.CopyTo(paths, 0);
+                paths[^1] = paths[^1].Parent.Combine(relativePath);
+                return new FullPath(Base, paths);
+            }
+        }
     }
 }

Wabbajack.Common/Paths/RelativePath.cs

@@ -139,6 +139,11 @@ namespace Wabbajack.Common
             return _path.StartsWith(s, StringComparison.OrdinalIgnoreCase);
         }
         
+        public bool EndsWith(string s)
+        {
+            return _path.EndsWith(s, StringComparison.OrdinalIgnoreCase);
+        }
+        
         public bool StartsWith(RelativePath s)
         {
             return _path.StartsWith(s._path, StringComparison.OrdinalIgnoreCase);

Wabbajack.ImageHashing/ImageState.cs

@@ -41,7 +41,7 @@ namespace Wabbajack.ImageHashing
             PerceptualHash.Write(bw);
         }
 
-        public static async Task<ImageState?> FromImageStream(Stream stream, Extension ext, bool takeStreamOwnership = true)
+        public static async Task<ImageState> FromImageStream(Stream stream, Extension ext, bool takeStreamOwnership = true)
         {
             await using var tf = new TempFile(ext);
             await tf.Path.WriteAllAsync(stream, takeStreamOwnership);
@@ -51,6 +51,9 @@ namespace Wabbajack.ImageHashing
         private static readonly Extension PNGExtension = new(".png");
         public static async Task<PHash> GetPHash(AbsolutePath path)
         {
+            if (!path.Exists)
+                throw new FileNotFoundException($"Can't hash non-existent file {path}");
+            
             await using var tmp = await TempFolder.Create();
             await ConvertImage(path, tmp.Dir, 512, 512, DXGI_FORMAT.R8G8B8A8_UNORM, PNGExtension);
             
@@ -75,12 +78,12 @@ namespace Wabbajack.ImageHashing
         public static async Task ConvertImage(Stream from, ImageState state, Extension ext, AbsolutePath to)
         {
             await using var tmpFile = await TempFolder.Create();
-            var inFile = to.FileName.RelativeTo(tmpFile.Dir).WithExtension(ext);
+            var inFile = to.FileName.RelativeTo(tmpFile.Dir);
             await inFile.WriteAllAsync(from);
             await ConvertImage(inFile, to.Parent, state.Width, state.Height, state.Format, ext);
         }
 
-        public static async Task<ImageState?> GetState(AbsolutePath path)
+        public static async Task<ImageState> GetState(AbsolutePath path)
         {
             var ph = new ProcessHelper
                 {
@@ -94,14 +97,7 @@ namespace Wabbajack.ImageHashing
                     .Select(p => p.Line)
                     .Where(p => p.Contains(" = "))
                     .Subscribe(l => lines.Push(l));
-                try
-                {
-                    await ph.Start();
-                }
-                catch (Exception ex)
-                {
-                    return null;
-                }
+                await ph.Start();
 
                 var data = lines.Select(l =>
                 {

Wabbajack.Lib/ACompiler.cs

@@ -21,14 +21,13 @@ namespace Wabbajack.Lib
 {
     public abstract class ACompiler : ABatchProcessor
     {
-        protected readonly Subject<(string, float)> _progressUpdates = new Subject<(string, float)>();
+        protected readonly Subject<(string, float)> _progressUpdates = new();
 
-        public List<IndexedArchive> IndexedArchives = new List<IndexedArchive>();
+        public List<IndexedArchive> IndexedArchives = new();
 
-        public Dictionary<Hash, IEnumerable<VirtualFile>> IndexedFiles =
-            new Dictionary<Hash, IEnumerable<VirtualFile>>();
+        public Dictionary<Hash, IEnumerable<VirtualFile>> IndexedFiles = new();
 
-        public ModList ModList = new ModList();
+        public ModList ModList = new();
         public AbsolutePath ModListImage;
         public bool ModlistIsNSFW;
 

Wabbajack.Lib/AuthorApi/Client.cs

@@ -44,12 +44,12 @@ namespace Wabbajack.Lib.AuthorApi
             IEnumerable<CDNFilePartDefinition> Blocks(AbsolutePath path)
             {
                 var size = path.Size;
-                for (long block = 0; block * Consts.UPLOADED_FILE_BLOCK_SIZE < size; block ++)
+                for (long block = 0; block * Consts.UploadedFileBlockSize < size; block ++)
                     yield return new CDNFilePartDefinition
                     {
                         Index = block,
-                        Size = Math.Min(Consts.UPLOADED_FILE_BLOCK_SIZE, size - block * Consts.UPLOADED_FILE_BLOCK_SIZE),
-                        Offset = block * Consts.UPLOADED_FILE_BLOCK_SIZE
+                        Size = Math.Min(Consts.UploadedFileBlockSize, size - block * Consts.UploadedFileBlockSize),
+                        Offset = block * Consts.UploadedFileBlockSize
                     };
             }
             

Wabbajack.Lib/CompilationSteps/DeconstructBSAs.cs

@@ -34,6 +34,7 @@ namespace Wabbajack.Lib.CompilationSteps
             _microstack = bsa => new List<ICompilationStep>
             {
                 new DirectMatch(_mo2Compiler),
+                new MatchSimilarTextures(_mo2Compiler),
                 new IncludePatches(_mo2Compiler, bsa),
                 new DropAll(_mo2Compiler)
             };
@@ -41,6 +42,7 @@ namespace Wabbajack.Lib.CompilationSteps
             _microstackWithInclude = bsa => new List<ICompilationStep>
             {
                 new DirectMatch(_mo2Compiler),
+                new MatchSimilarTextures(_mo2Compiler),
                 new IncludePatches(_mo2Compiler, bsa),
                 new IncludeAll(_mo2Compiler)
             };

Wabbajack.Lib/CompilationSteps/MatchSimilarTextures.cs

@@ -1,5 +1,8 @@
-using System.Linq;
+using System;
+using System.IO;
+using System.Linq;
 using System.Threading.Tasks;
+using Microsoft.VisualBasic.Logging;
 using Wabbajack.Common;
 using Wabbajack.ImageHashing;
 using Wabbajack.VirtualFileSystem;
@@ -17,24 +20,63 @@ namespace Wabbajack.Lib.CompilationSteps
                 .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.Path.Extension == DDS && source.File.ImageState != null)
+            if (source.File.Name.FileName.Extension == DDS && source.File.ImageState != null)
             {
-                var found = _byName[source.Path.FileNameWithoutExtension]
+                (float Similarity, VirtualFile File) found = _byName[source.Path.FileNameWithoutExtension]
                     .Select(f => (f.ImageState.PerceptualHash.Similarity(source.File.ImageState.PerceptualHash), f))
-                    .Where(f => f.Item1 >= 0.90f)
+                    .Select(f =>
+                    {
+                        Utils.Log($"{f.f.Name.FileName} similar {f.Item1}");
+                        return f;
+                    })
                     .OrderByDescending(f => f.Item1)
                     .FirstOrDefault();
 
-                if (found == default) return null;
+                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(new RelativePath(mainFileName))
+                        where mainFile != null
+                        from mainMatch in _byName[mainFile.FullPath.FileName.FileNameWithoutExtension]
+                        where mainMatch.ImageState != null
+                        where mainFile.ImageState != null
+                            let similarity = mainFile.ImageState.PerceptualHash.Similarity(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);
+
+                    var foundRec = r.FirstOrDefault();
+                    if (foundRec == default)
+                    {
+                        return null;
+                    }
+
+                    found = (foundRec.similarity, foundRec.foundFile);
+                }
 
                 var rv = source.EvolveTo<TransformedTexture>();
-                rv.ArchiveHashPath = found.f.MakeRelativePaths();
-                rv.ImageState = found.f.ImageState;
+                rv.ArchiveHashPath = found.File.MakeRelativePaths();
+                rv.ImageState = found.File.ImageState;
                 
                 return rv;
             }

Wabbajack.Server.Test/AuthoredFilesTests.cs

@@ -28,7 +28,7 @@ namespace Wabbajack.BuildServer.Test
             var toDelete = await cleanup.FindFilesToDelete();
             
             await using var file = new TempFile();
-            await file.Path.WriteAllBytesAsync(RandomData(Consts.UPLOADED_FILE_BLOCK_SIZE * 4 + Consts.UPLOADED_FILE_BLOCK_SIZE / 3));
+            await file.Path.WriteAllBytesAsync(RandomData(Consts.UploadedFileBlockSize * 4 + Consts.UploadedFileBlockSize / 3));
             var originalHash = await file.Path.FileHashAsync();
 
             var client = await Client.Create(Fixture.APIKey);

Wabbajack.Test/SanityTests.cs

@@ -387,15 +387,18 @@ namespace Wabbajack.Test
             {
                 var originalDDS = await ImageState.GetState(nativeFile);
                 await ImageState.ConvertImage(nativeFile, recompressedFile.Parent, originalDDS.Width, originalDDS.Height, DXGI_FORMAT.BC7_UNORM, recompressedFile.Extension);
-                await ImageState.ConvertImage(nativeFile, resizedFile.Parent, 128, 128, DXGI_FORMAT.BC7_UNORM, resizedFile.Extension);
+                await ImageState.ConvertImage(nativeFile, resizedFile.Parent, 1024, 1024, DXGI_FORMAT.BC7_UNORM, resizedFile.Extension);
             }
 
             await utils.Configure();
             
-            await CompileAndInstall(profile, true);
+            var compilerData = await CompileAndInstall(profile, true);
             await utils.VerifyInstalledFile(mod, @"native\whitestagbody.dds");
-            Assert.True(0.99f <=(await ImageState.GetState(recompressedFile)).PerceptualHash.Similarity(await ImageState.GetPHash(utils.InstalledPath(mod, @"recompressed\whitestagbody.dds"))));
-            Assert.True(0.98f <=(await ImageState.GetState(resizedFile)).PerceptualHash.Similarity(await ImageState.GetPHash(utils.InstalledPath(mod, @"resized\whitestagbody.dds"))));
+
+            var directies = compilerData.Directives;
+            
+            Assert.True(0.99f <= (await ImageState.GetState(recompressedFile)).PerceptualHash.Similarity(await ImageState.GetPHash(utils.InstalledPath(mod, @"recompressed\whitestagbody.dds"))));
+            Assert.True(0.98f <= (await ImageState.GetState(resizedFile)).PerceptualHash.Similarity(await ImageState.GetPHash(utils.InstalledPath(mod, @"resized\whitestagbody.dds"))));
         }
         
         [Fact]
@@ -697,7 +700,7 @@ namespace Wabbajack.Test
             var gameFolder = Game.SkyrimSpecialEdition.MetaData().GameLocation();
             await gameFolder.Combine("SkyrimSE.exe").CopyToAsync(utils.SourcePath.Combine("SkyrimSE.exe"));
 
-            var some_dds = utils.SourcePath.Combine("some_file.dds");
+            var some_dds = utils.SourcePath.Combine("some_file.txx");
             await some_dds.WriteAllBytesAsync(utils.RandomData());
 
             var blerg = utils.SourcePath.Combine("file1.blerg");
@@ -726,7 +729,7 @@ namespace Wabbajack.Test
 
             await CompileAndInstall(settingsPath, true);
             
-            Assert.Equal(await some_dds.FileHashAsync(), await utils.InstallPath.Combine("some_file.dds").FileHashAsync());
+            Assert.Equal(await some_dds.FileHashAsync(), await utils.InstallPath.Combine("some_file.txx").FileHashAsync());
             Assert.Equal(await gameFolder.Combine("SkyrimSE.exe").FileHashAsync(), 
                 await utils.InstallPath.Combine("SkyrimSE.exe").FileHashAsync());
         }

Wabbajack.VirtualFileSystem/IndexedVirtualFile.cs

@@ -46,10 +46,7 @@ namespace Wabbajack.VirtualFileSystem
         {
             using var cs = new GZipStream(s, CompressionLevel.Optimal , true);
             using var bw = new BinaryWriter(cs, Encoding.UTF8, true);
-            bw.Write(Size);
-            bw.Write(Children.Count);
-            foreach (var file in Children)
-                file.Write(bw);
+            Write(bw);
         }
 
         private static IndexedVirtualFile Read(BinaryReader br)
@@ -80,18 +77,7 @@ namespace Wabbajack.VirtualFileSystem
         {
             using var cs = new GZipStream(s, CompressionMode.Decompress, true);
             using var br = new BinaryReader(cs);
-            var ivf = new IndexedVirtualFile
-            {
-                Size = br.ReadInt64(),
-            };
-            var lst = new List<IndexedVirtualFile>();
-            ivf.Children = lst;
-            var count = br.ReadInt32();
-            for (int x = 0; x < count; x++)
-            {
-                lst.Add(Read(br));
-            }
-            return ivf;
+            return Read(br);
         }
     }
 }

Wabbajack.VirtualFileSystem/VirtualFile.cs

@@ -17,7 +17,7 @@ namespace Wabbajack.VirtualFileSystem
 {
     public class VirtualFile
     {
-        private static AbsolutePath DBLocation = Consts.LocalAppDataPath.Combine("GlobalVFSCache2.sqlite");
+        private static AbsolutePath DBLocation = Consts.LocalAppDataPath.Combine("GlobalVFSCache3.sqlite");
         private static string _connectionString;
         private static SQLiteConnection _conn;
 
@@ -227,7 +227,10 @@ namespace Wabbajack.VirtualFileSystem
             };
 
             if (Consts.TextureExtensions.Contains(relPath.FileName.Extension))
+            {
                 self.ImageState = await ImageState.FromImageStream(stream, relPath.FileName.Extension, false);
+                stream.Position = 0;
+            }
 
             self.FillFullPath(depth);
             
@@ -267,7 +270,11 @@ namespace Wabbajack.VirtualFileSystem
         private static async Task WriteToCache(VirtualFile self)
         {
             await using var ms = new MemoryStream();
-            self.ToIndexedVirtualFile().Write(ms);
+            var ivf = self.ToIndexedVirtualFile();
+            // Top level path gets renamed when read, we don't want the absolute path
+            // here else the reader will blow up when it tries to convert the value
+            ivf.Name = (RelativePath)"not/applicable";
+            ivf.Write(ms);
             ms.Position = 0;
             await InsertIntoVFSCache(self.Hash, ms);
         }
@@ -436,6 +443,12 @@ namespace Wabbajack.VirtualFileSystem
             var path = new HashRelativePath(FilesInFullPath.First().Hash, paths);
             return path;
         }
+
+        public VirtualFile InSameFolder(RelativePath relativePath)
+        {
+            var newPath = FullPath.InSameFolder(relativePath);
+            return Context.Index.ByFullPath.TryGetValue(newPath, out var found) ? found : null;
+        }
     }
 
     public class ExtendedHashes