Full round trip test passes

This commit is contained in:
Timothy Baldridge 2021-06-17 17:09:03 -06:00
parent 6706d5bfb6
commit 528ef51d40
12 changed files with 244 additions and 36 deletions

View File

@ -128,7 +128,7 @@ namespace Wabbajack.Common
return Hash.FromBase64((string)reader.Value!);
}
}
public class RelativePathConverter : JsonConverter<RelativePath>
{
public override void WriteJson(JsonWriter writer, RelativePath value, JsonSerializer serializer)

View File

@ -1,8 +1,10 @@
using System;
using System.Collections.Generic;
using System.IO;
using System.Numerics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using DirectXTexNet;
using Shipwreck.Phash;
using Shipwreck.Phash.Imaging;
@ -10,7 +12,7 @@ using Wabbajack.Common;
namespace Wabbajack.ImageHashing
{
public class DDSImage
public class DDSImage : IDisposable
{
private DDSImage(ScratchImage img, TexMetadata metadata, Extension ext)
{
@ -20,6 +22,7 @@ namespace Wabbajack.ImageHashing
}
private static Extension DDSExtension = new(".dds");
private static Extension TGAExtension = new(".tga");
private ScratchImage _image;
private TexMetadata _metaData;
@ -32,7 +35,7 @@ namespace Wabbajack.ImageHashing
return new DDSImage(img, img.GetMetadata(), new Extension(".dds"));
}
public static DDSImage FromDDSMemory(byte[] data)
{
unsafe
@ -56,6 +59,17 @@ namespace Wabbajack.ImageHashing
}
}
}
public static async Task<DDSImage> FromStream(Stream stream, IPath arg1Name)
{
var data = await stream.ReadAllAsync();
if (arg1Name.FileName.Extension == DDSExtension)
return FromDDSMemory(data);
if (arg1Name.FileName.Extension == TGAExtension)
return FromTGAMemory(data);
throw new NotImplementedException("Only DDS and TGA files supported");
}
public void Dispose()
{
@ -66,9 +80,34 @@ namespace Wabbajack.ImageHashing
public int Width => _metaData.Width;
public int Height => _metaData.Height;
public void Resize(int width, int height)
// Destructively resize a Image
public void ResizeRecompressAndSave(int width, int height, DXGI_FORMAT newFormat, AbsolutePath dest)
{
ScratchImage? resized = default;
try
{
// First we resize the image, so that changes due to image scaling matter less in the final hash
if (CompressedTypes.Contains(_metaData.Format))
{
using var decompressed = _image.Decompress(DXGI_FORMAT.UNKNOWN);
resized = decompressed.Resize(width, height, TEX_FILTER_FLAGS.DEFAULT);
}
else
{
resized = _image.Resize(width, height, TEX_FILTER_FLAGS.DEFAULT);
}
using var compressed = resized.Compress(newFormat, TEX_COMPRESS_FLAGS.BC7_QUICK, 0.5f);
if (dest.Extension == new Extension(".dds"))
{
compressed.SaveToDDSFile(DDS_FLAGS.NONE, dest.ToString());
}
}
finally
{
resized?.Dispose();
}
}
private static HashSet<DXGI_FORMAT> CompressedTypes = new HashSet<DXGI_FORMAT>()
@ -104,6 +143,7 @@ namespace Wabbajack.ImageHashing
{
Width = _metaData.Width,
Height = _metaData.Height,
Format = _metaData.Format,
PerceptualHash = PerceptionHash()
};
}
@ -159,6 +199,10 @@ namespace Wabbajack.ImageHashing
resized?.Dispose();
}
}
public void ResizeRecompressAndSave(ImageState state, AbsolutePath dest)
{
ResizeRecompressAndSave(state.Width, state.Height, state.Format, dest);
}
}
}

View File

@ -1,9 +1,13 @@
using System.IO;
using System;
using System.IO;
using System.Threading.Tasks;
using DirectXTexNet;
using Wabbajack.Common;
using Wabbajack.Common.Serialization.Json;
namespace Wabbajack.ImageHashing
{
[JsonName("ImageState")]
public class ImageState
{
public int Width { get; set; }
@ -29,5 +33,34 @@ namespace Wabbajack.ImageHashing
bw.Write((byte)Format);
PerceptualHash.Write(bw);
}
public static async Task<ImageState> FromImageStream(Stream stream, Extension ext, bool takeStreamOwnership = true)
{
var ms = new MemoryStream();
await stream.CopyToAsync(ms);
if (takeStreamOwnership) await stream.DisposeAsync();
DDSImage? img = default;
try
{
if (ext == new Extension(".dds"))
img = DDSImage.FromDDSMemory(ms.GetBuffer());
else if (ext == new Extension(".tga"))
{
img = DDSImage.FromTGAMemory(ms.GetBuffer());
}
else
{
throw new NotImplementedException("Only DDS and TGA files supported by PHash");
}
return img.ImageState();
}
finally
{
img?.Dispose();
}
}
}
}

View File

@ -2,31 +2,34 @@
using System.Data;
using System.IO;
using System.Threading.Tasks;
using Newtonsoft.Json;
using Shipwreck.Phash;
using Wabbajack.Common;
namespace Wabbajack.ImageHashing
{
[JsonConverter(typeof(PHashJsonConverter))]
public struct PHash
{
private const int SIZE = 40;
private readonly byte[] _data;
private readonly int _hash;
public byte[] Data { get; }
private PHash(byte[] data)
{
_data = data;
if (_data.Length != SIZE)
Data = data;
if (Data.Length != SIZE)
throw new DataException();
long h = 0;
h |= _data[0];
h |= Data[0];
h <<= 8;
h |= _data[1];
h |= Data[1];
h <<= 8;
h |= _data[2];
h |= Data[2];
h <<= 8;
h |= _data[3];
h |= Data[3];
h <<= 8;
_hash = (int)h;
}
@ -49,7 +52,7 @@ namespace Wabbajack.ImageHashing
if (_hash == 0)
br.Write(new byte[SIZE]);
else
br.Write(_data);
br.Write(Data);
}
public static PHash FromDigest(Digest digest)
@ -59,24 +62,24 @@ namespace Wabbajack.ImageHashing
public float Similarity(PHash other)
{
return ImagePhash.GetCrossCorrelation(this._data, other._data);
return ImagePhash.GetCrossCorrelation(this.Data, other.Data);
}
public override string ToString()
{
return _data.ToBase64();
return Data.ToBase64();
}
public override int GetHashCode()
{
long h = 0;
h |= _data[0];
h |= Data[0];
h <<= 8;
h |= _data[1];
h |= Data[1];
h <<= 8;
h |= _data[2];
h |= Data[2];
h <<= 8;
h |= _data[3];
h |= Data[3];
h <<= 8;
return (int)h;
}
@ -100,5 +103,27 @@ namespace Wabbajack.ImageHashing
}
return img.PerceptionHash();
}
public static async Task<PHash> FromFile(AbsolutePath path)
{
await using var s = await path.OpenRead();
return await FromStream(s, path.Extension);
}
}
public class PHashJsonConverter : JsonConverter<PHash>
{
public override void WriteJson(JsonWriter writer, PHash value, JsonSerializer serializer)
{
writer.WriteValue(value.Data.ToBase64());
}
public override PHash ReadJson(JsonReader reader, Type objectType, PHash existingValue, bool hasExistingValue,
JsonSerializer serializer)
{
return PHash.FromBase64((string)reader.Value!);
}
}
}

View File

@ -6,6 +6,7 @@ using System.Linq;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using Wabbajack.Common;
using Wabbajack.ImageHashing;
using Wabbajack.Lib.Downloaders;
using Wabbajack.Lib.Validation;
using Wabbajack.VirtualFileSystem;
@ -153,6 +154,16 @@ namespace Wabbajack.Lib
}
break;
case TransformedTexture tt:
{
await using var s = await sf.GetStream();
using var img = await DDSImage.FromStream(s, vf.Name);
img.ResizeRecompressAndSave(tt.ImageState, directive.Directive.To.RelativeTo(OutputFolder));
}
break;
case FromArchive _:
if (grouped[vf].Count() == 1)
{

View File

@ -0,0 +1,44 @@
using System.Linq;
using System.Threading.Tasks;
using Wabbajack.Common;
using Wabbajack.ImageHashing;
using Wabbajack.VirtualFileSystem;
namespace Wabbajack.Lib.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)
.ToLookup(f => f.Name.FileName.FileNameWithoutExtension);
}
private static Extension DDS = new(".dds");
public override async ValueTask<Directive?> Run(RawSourceFile source)
{
if (source.Path.Extension == DDS)
{
var found = _byName[source.Path.FileNameWithoutExtension]
.Select(f => (f.ImageState.PerceptualHash.Similarity(source.File.ImageState.PerceptualHash), f))
.Where(f => f.Item1 >= 0.90f)
.OrderByDescending(f => f.Item1)
.FirstOrDefault();
if (found == default) return null;
var rv = source.EvolveTo<TransformedTexture>();
rv.ArchiveHashPath = found.f.MakeRelativePaths();
rv.ImageState = found.f.ImageState;
return rv;
}
return null;
}
}
}

View File

@ -6,6 +6,7 @@ using Compression.BSA;
using Newtonsoft.Json;
using Wabbajack.Common;
using Wabbajack.Common.Serialization.Json;
using Wabbajack.ImageHashing;
using Wabbajack.Lib.Downloaders;
using Wabbajack.VirtualFileSystem;
@ -37,10 +38,7 @@ namespace Wabbajack.Lib
public T EvolveTo<T>() where T : Directive, new()
{
var v = new T();
v.To = Path;
v.Hash = File.Hash;
v.Size = File.Size;
var v = new T {To = Path, Hash = File.Hash, Size = File.Size};
return v;
}
}
@ -252,6 +250,15 @@ namespace Wabbajack.Lib
[JsonIgnore]
public VirtualFile[] Choices { get; set; } = { };
}
[JsonName("TransformedTexture")]
public class TransformedTexture : FromArchive
{
/// <summary>
/// The file to apply to the source file to patch it
/// </summary>
public ImageState ImageState { get; set; } = new();
}
[JsonName("SourcePatch")]
public class SourcePatch

View File

@ -456,6 +456,7 @@ namespace Wabbajack.Lib
new IgnoreEndsWith(this, ".log"),
new DeconstructBSAs(
this), // Deconstruct BSAs before building patches so we don't generate massive patch files
new MatchSimilarTextures(this),
new IncludePatches(this),
new IncludeDummyESPs(this),

View File

@ -5,12 +5,14 @@ using System.Linq;
using System.Threading.Tasks;
using Compression.BSA;
using Wabbajack.Common;
using Wabbajack.ImageHashing;
using Wabbajack.Lib;
using Wabbajack.Lib.CompilationSteps;
using Wabbajack.Lib.CompilationSteps.CompilationErrors;
using Xunit;
using Xunit.Abstractions;
using Xunit.Sdk;
using DXGI_FORMAT = DirectXTexNet.DXGI_FORMAT;
using File = Alphaleonis.Win32.Filesystem.File;
using Path = Alphaleonis.Win32.Filesystem.Path;
@ -364,6 +366,41 @@ namespace Wabbajack.Test
}
[Fact]
public async Task CanRecompressAndResizeDDSImages()
{
var profile = utils.AddProfile();
var mod = await utils.AddMod();
var nativeFile = await utils.AddModFile(mod, @"native\whitestagbody.dds", 0);
var recompressedFile = await utils.AddModFile(mod, @"recompressed\whitestagbody.dds", 0);
var resizedFile = await utils.AddModFile(mod, @"resized\whitestagbody.dds", 0);
var gameBSA = Game.SkyrimSpecialEdition.MetaData().GameLocation().Combine(@"Data\Skyrim - Textures1.bsa");
var bsa = await BSADispatch.OpenRead(gameBSA);
var ddsExtension = new Extension(".dds");
var firstFile = bsa.Files.First(f => f.Path.Extension == ddsExtension);
await using (var nf = await nativeFile.OpenWrite())
{
await firstFile.CopyDataTo(nf);
}
{
using var originalDDS = DDSImage.FromFile(nativeFile);
originalDDS.ResizeRecompressAndSave(originalDDS.Width, originalDDS.Height, DXGI_FORMAT.BC7_UNORM, recompressedFile);
originalDDS.ResizeRecompressAndSave(128, 128, DXGI_FORMAT.BC7_UNORM, resizedFile);
}
await utils.Configure();
await CompileAndInstall(profile, true);
await utils.VerifyInstalledFile(mod, @"native\whitestagbody.dds");
Assert.Equal(0.9999227f, (await PHash.FromFile(recompressedFile)).Similarity(await PHash.FromFile(utils.InstalledPath(mod, @"recompressed\whitestagbody.dds"))));
Assert.Equal(0.98703325f, (await PHash.FromFile(resizedFile)).Similarity(await PHash.FromFile(utils.InstalledPath(mod, @"resized\whitestagbody.dds"))));
}
[Fact]
public async Task CanNoMatchIncludeFilesFromBSAs()
{

View File

@ -95,14 +95,17 @@ namespace Wabbajack.Test
/// </summary>
/// <param name="mod_name"></param>
/// <param name="path"></param>
/// <param name="random_fill"></param>
/// <param name="randomFill"></param>
/// <returns></returns>
public async Task<AbsolutePath> AddModFile(string mod_name, string path, int random_fill=128)
public async Task<AbsolutePath> AddModFile(string mod_name, string path, int randomFill=128)
{
var full_path = ModsPath.Combine(mod_name, path);
full_path.Parent.CreateDirectory();
await GenerateRandomFileData(full_path, random_fill);
return full_path;
var fullPath = ModsPath.Combine(mod_name, path);
fullPath.Parent.CreateDirectory();
if (randomFill != 0)
await GenerateRandomFileData(fullPath, randomFill);
return fullPath;
}
public async Task GenerateRandomFileData(AbsolutePath full_path, int random_fill)
@ -198,6 +201,9 @@ namespace Wabbajack.Test
Assert.True(false, $"Index {x} of {mod}\\{file} are not the same");
}
}
public AbsolutePath InstalledPath(string mod, string file) =>
InstallPath.Combine((string)Consts.MO2ModFolderName, mod, file);
public async Task VerifyInstalledGameFile(string file)
{

View File

@ -19,7 +19,7 @@ namespace Wabbajack.VirtualFileSystem
public IPath Name { get; set; }
public Hash Hash { get; set; }
public ImageState? ImageState { get; set; }
public ImageState ImageState { get; set; }
public long Size { get; set; }
public List<IndexedVirtualFile> Children { get; set; } = new();

View File

@ -49,8 +49,7 @@ namespace Wabbajack.VirtualFileSystem
public FullPath FullPath { get; private set; }
public Hash Hash { get; internal set; }
public PHash PerceptualHash { get; internal set; }
public ImageState ImageState { get; internal set; }
public ExtendedHashes ExtendedHashes { get; set; }
public long Size { get; internal set; }
@ -148,7 +147,8 @@ namespace Wabbajack.VirtualFileSystem
Size = file.Size,
LastModified = extractedFile.LastModifiedUtc.AsUnixTime(),
LastAnalyzed = DateTime.Now.AsUnixTime(),
Hash = file.Hash
Hash = file.Hash,
ImageState = file.ImageState
};
vself.FillFullPath();
@ -183,7 +183,7 @@ namespace Wabbajack.VirtualFileSystem
return new()
{
Hash = Hash,
PerceptualHash = PerceptualHash,
ImageState = ImageState,
Name = Name,
Children = Children.Select(c => c.ToIndexedVirtualFile()).ToList(),
Size = Size
@ -227,7 +227,7 @@ namespace Wabbajack.VirtualFileSystem
};
if (Consts.TextureExtensions.Contains(relPath.FileName.Extension))
self.PerceptualHash = await PHash.FromStream(stream, relPath.FileName.Extension, false);
self.ImageState = await ImageHashing.ImageState.FromImageStream(stream, relPath.FileName.Extension, false);
self.FillFullPath(depth);