diff --git a/Wabbajack.Common/Consts.cs b/Wabbajack.Common/Consts.cs index 9c40ad32..065ef487 100644 --- a/Wabbajack.Common/Consts.cs +++ b/Wabbajack.Common/Consts.cs @@ -172,5 +172,6 @@ namespace Wabbajack.Common public static long UPLOADED_FILE_BLOCK_SIZE = (long)1024 * 1024 * 2; + public static HashSet TextureExtensions = new() {new Extension(".dds"), new Extension(".tga")}; } } diff --git a/Wabbajack.ImageHashing.Test/ImageLoadingTests.cs b/Wabbajack.ImageHashing.Test/ImageLoadingTests.cs index 4c89a910..849f7b10 100644 --- a/Wabbajack.ImageHashing.Test/ImageLoadingTests.cs +++ b/Wabbajack.ImageHashing.Test/ImageLoadingTests.cs @@ -7,9 +7,52 @@ namespace Wabbajack.ImageHashing.Test public class ImageLoadingTests { [Fact] - public async Task CanLoadDDSFiles() + public async Task CanLoadAndCompareDDSImages() { - var file = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5.dds")); + var file1 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5.dds")); + var hash1 = file1.PerceptionHash(); + + var file2 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5.dds")); + var hash2 = file2.PerceptionHash(); + + Assert.Equal(1, hash1.Similarity(hash2)); + } + + [Fact] + public async Task CanLoadAndCompareResizedImage() + { + var file1 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5.dds")); + var hash1 = file1.PerceptionHash(); + + var file2 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5-small-bc7.dds")); + var hash2 = file2.PerceptionHash(); + + Assert.Equal(0.956666886806488, hash1.Similarity(hash2)); + } + + + [Fact] + public async Task CanLoadAndCompareResizedVFlipImage() + { + var file1 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5.dds")); + var hash1 = file1.PerceptionHash(); + + var file2 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5-small-bc7-vflip.dds")); + var hash2 = file2.PerceptionHash(); + + Assert.Equal(0.2465425431728363, hash1.Similarity(hash2)); + } + + [Fact] + public async Task CanLoadAndCompareRecompressedImage() + { + var file1 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5.dds")); + var hash1 = file1.PerceptionHash(); + + var file2 = DDSImage.FromFile(AbsolutePath.EntryPoint.Combine("Resources", "test-dxt5-recompressed.dds")); + var hash2 = file2.PerceptionHash(); + + Assert.Equal(0.9999724626541138, hash1.Similarity(hash2)); } } } diff --git a/Wabbajack.ImageHashing.Test/Resources/test-dxt5-recompressed.dds b/Wabbajack.ImageHashing.Test/Resources/test-dxt5-recompressed.dds new file mode 100644 index 00000000..1b90c6e1 Binary files /dev/null and b/Wabbajack.ImageHashing.Test/Resources/test-dxt5-recompressed.dds differ diff --git a/Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7-vflip.dds b/Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7-vflip.dds new file mode 100644 index 00000000..13942ad9 Binary files /dev/null and b/Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7-vflip.dds differ diff --git a/Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7.dds b/Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7.dds new file mode 100644 index 00000000..911720d4 Binary files /dev/null and b/Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7.dds differ diff --git a/Wabbajack.ImageHashing.Test/Wabbajack.ImageHashing.Test.csproj b/Wabbajack.ImageHashing.Test/Wabbajack.ImageHashing.Test.csproj index e264262f..d7e7f80f 100644 --- a/Wabbajack.ImageHashing.Test/Wabbajack.ImageHashing.Test.csproj +++ b/Wabbajack.ImageHashing.Test/Wabbajack.ImageHashing.Test.csproj @@ -4,6 +4,8 @@ net5.0-windows false + + enable @@ -27,6 +29,15 @@ PreserveNewest + + PreserveNewest + + + PreserveNewest + + + PreserveNewest + diff --git a/Wabbajack.ImageHashing/DDSImage.cs b/Wabbajack.ImageHashing/DDSImage.cs index d0ad5603..6f8fd305 100644 --- a/Wabbajack.ImageHashing/DDSImage.cs +++ b/Wabbajack.ImageHashing/DDSImage.cs @@ -1,14 +1,21 @@ using System; +using System.Collections.Generic; +using System.Numerics; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; using DirectXTexNet; +using Shipwreck.Phash; +using Shipwreck.Phash.Imaging; using Wabbajack.Common; namespace Wabbajack.ImageHashing { - public class DDSImage : IImage + public class DDSImage { - private DDSImage() + private DDSImage(ScratchImage img, TexMetadata metadata) { - + _image = img; + _metaData = metadata; } private static Extension DDSExtension = new(".dds"); @@ -22,7 +29,31 @@ namespace Wabbajack.ImageHashing var img = TexHelper.Instance.LoadFromDDSFile(file.ToString(), DDS_FLAGS.NONE); - return new DDSImage() {_image = img, _metaData = img.GetMetadata()}; + return new DDSImage(img, img.GetMetadata()); + } + + public static DDSImage FromDDSMemory(byte[] data) + { + unsafe + { + fixed (byte* ptr = data) + { + var img = TexHelper.Instance.LoadFromDDSMemory((IntPtr)ptr, data.Length, DDS_FLAGS.NONE); + return new DDSImage(img, img.GetMetadata()); + } + } + } + + public static DDSImage FromTGAMemory(byte[] data) + { + unsafe + { + fixed (byte* ptr = data) + { + var img = TexHelper.Instance.LoadFromTGAMemory((IntPtr)ptr, data.Length); + return new DDSImage(img, img.GetMetadata()); + } + } } public void Dispose() @@ -33,7 +64,89 @@ namespace Wabbajack.ImageHashing public int Width => _metaData.Width; public int Height => _metaData.Height; - public GPUCompressionLevel CompressionLevel => GPUCompressionLevel.Uncompressed; - public IImageState State { get; } + + public void Resize(int width, int height) + { + + } + + private static HashSet CompressedTypes = new HashSet() + { + DXGI_FORMAT.BC1_TYPELESS, + DXGI_FORMAT.BC1_UNORM, + DXGI_FORMAT.BC1_UNORM_SRGB, + DXGI_FORMAT.BC2_TYPELESS, + DXGI_FORMAT.BC2_UNORM, + DXGI_FORMAT.BC2_UNORM_SRGB, + DXGI_FORMAT.BC3_TYPELESS, + DXGI_FORMAT.BC3_UNORM, + DXGI_FORMAT.BC3_UNORM_SRGB, + DXGI_FORMAT.BC4_TYPELESS, + DXGI_FORMAT.BC4_UNORM, + DXGI_FORMAT.BC4_SNORM, + DXGI_FORMAT.BC5_TYPELESS, + DXGI_FORMAT.BC5_UNORM, + DXGI_FORMAT.BC5_SNORM, + DXGI_FORMAT.BC6H_TYPELESS, + DXGI_FORMAT.BC6H_UF16, + DXGI_FORMAT.BC6H_SF16, + DXGI_FORMAT.BC7_TYPELESS, + DXGI_FORMAT.BC7_UNORM, + DXGI_FORMAT.BC7_UNORM_SRGB, + }; + + public PHash PerceptionHash() + { + ScratchImage? resized = default; + try + { + if (CompressedTypes.Contains(_metaData.Format)) + { + using var decompressed = _image.Decompress(DXGI_FORMAT.UNKNOWN); + resized = decompressed.Resize(512, 512, TEX_FILTER_FLAGS.DEFAULT); + } + else + { + resized = _image.Resize(512, 512, TEX_FILTER_FLAGS.DEFAULT); + } + + var data = new List<(int, int)>(); + var image = new byte[512 * 512]; + + unsafe void EvaluatePixels(IntPtr pixels, IntPtr width, IntPtr line) + { + float* ptr = (float*)pixels.ToPointer(); + + int widthV = width.ToInt32(); + if (widthV != 512) return; + + var y = line.ToInt32(); + data.Add((widthV, y)); + + for (int i = 0; i < widthV; i++) + { + var r = ptr[0] * 0.229f; + var g = ptr[1] * 0.587f; + var b = ptr[2] * 0.114f; + + var combined = (r + g + b) * 255.0f; + + image[(y * widthV) + i] = (byte)combined; + ptr += 4; + } + + } + + resized.EvaluateImage(EvaluatePixels); + + var digest = ImagePhash.ComputeDigest(new ByteImage(512, 512, image)); + return PHash.FromDigest(digest); + } + finally + { + resized?.Dispose(); + } + } + } } diff --git a/Wabbajack.ImageHashing/IImage.cs b/Wabbajack.ImageHashing/IImage.cs index 00d1aaf2..6973124a 100644 --- a/Wabbajack.ImageHashing/IImage.cs +++ b/Wabbajack.ImageHashing/IImage.cs @@ -2,23 +2,5 @@ namespace Wabbajack.ImageHashing { - public interface IImage : IDisposable - { - public int Width { get; } - public int Height { get; } - public GPUCompressionLevel CompressionLevel { get; } - public IImageState State { get; } - } - public interface IImageState - { - - } - - public enum GPUCompressionLevel : int - { - Uncompressed = 0, // The Image is uncompressed on the GPU - Old = 1, // The Image is compressed in a poor (old) format on the GPU - New = 2 // The Image is compressed in a newer format (like BC7) on the GPU - } } diff --git a/Wabbajack.ImageHashing/PHash.cs b/Wabbajack.ImageHashing/PHash.cs new file mode 100644 index 00000000..dab74b97 --- /dev/null +++ b/Wabbajack.ImageHashing/PHash.cs @@ -0,0 +1,89 @@ +using System; +using System.Data; +using System.IO; +using System.Threading.Tasks; +using Shipwreck.Phash; +using Wabbajack.Common; + +namespace Wabbajack.ImageHashing +{ + public struct PHash + { + private const int SIZE = 40; + private readonly byte[] _data; + + private PHash(byte[]? data) + { + _data = data ?? new byte[SIZE]; + if (_data.Length != SIZE) + throw new DataException(); + } + + public static PHash FromBase64(string base64) + { + var data = base64.FromBase64(); + if (data.Length != SIZE) + throw new DataException(); + return new PHash(data); + } + + public static PHash Read(BinaryReader br) + { + return new (br.ReadBytes(SIZE)); + } + + public void Write(BinaryWriter br) + { + br.Write((_data?.Length ?? 0) == 0 ? new byte[SIZE] : _data); + } + + public static PHash FromDigest(Digest digest) + { + return new(digest.Coefficients); + } + + public float Similarity(PHash other) + { + return ImagePhash.GetCrossCorrelation(this._data, other._data); + } + + public override string ToString() + { + return _data.ToBase64(); + } + + public override int GetHashCode() + { + long h = 0; + h |= _data[0]; + h <<= 8; + h |= _data[1]; + h <<= 8; + h |= _data[2]; + h <<= 8; + h |= _data[3]; + h <<= 8; + return (int)h; + } + + public static async Task FromStream(Stream stream, Extension ext, bool takeStreamOwnership = true) + { + var ms = new MemoryStream(); + await stream.CopyToAsync(ms); + if (takeStreamOwnership) await stream.DisposeAsync(); + + DDSImage img; + 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.PerceptionHash(); + } + } +} diff --git a/Wabbajack.ImageHashing/Wabbajack.ImageHashing.csproj b/Wabbajack.ImageHashing/Wabbajack.ImageHashing.csproj index d11721ec..26a87344 100644 --- a/Wabbajack.ImageHashing/Wabbajack.ImageHashing.csproj +++ b/Wabbajack.ImageHashing/Wabbajack.ImageHashing.csproj @@ -2,11 +2,13 @@ net5.0-windows + true + enable - + diff --git a/Wabbajack.VirtualFileSystem/IndexedVirtualFile.cs b/Wabbajack.VirtualFileSystem/IndexedVirtualFile.cs index cca6a44f..fa01224c 100644 --- a/Wabbajack.VirtualFileSystem/IndexedVirtualFile.cs +++ b/Wabbajack.VirtualFileSystem/IndexedVirtualFile.cs @@ -5,6 +5,7 @@ using System.Text; using System.Threading.Tasks; using Wabbajack.Common; using Wabbajack.Common.Serialization.Json; +using Wabbajack.ImageHashing; namespace Wabbajack.VirtualFileSystem { @@ -16,13 +17,16 @@ namespace Wabbajack.VirtualFileSystem { public IPath Name { get; set; } public Hash Hash { get; set; } + + public PHash PerceptualHash { get; set; } public long Size { get; set; } - public List Children { get; set; } = new List(); + public List Children { get; set; } = new(); private void Write(BinaryWriter bw) { bw.Write(Name.ToString()); bw.Write((ulong)Hash); + PerceptualHash.Write(bw); bw.Write(Size); bw.Write(Children.Count); foreach (var file in Children) @@ -44,6 +48,7 @@ namespace Wabbajack.VirtualFileSystem { Name = (RelativePath)br.ReadString(), Hash = Hash.FromULong(br.ReadUInt64()), + PerceptualHash = PHash.Read(br), Size = br.ReadInt64(), }; var lst = new List(); diff --git a/Wabbajack.VirtualFileSystem/VirtualFile.cs b/Wabbajack.VirtualFileSystem/VirtualFile.cs index ff92f29e..34fd5b1f 100644 --- a/Wabbajack.VirtualFileSystem/VirtualFile.cs +++ b/Wabbajack.VirtualFileSystem/VirtualFile.cs @@ -10,6 +10,8 @@ using System.Threading.Tasks; using ICSharpCode.SharpZipLib.Zip.Compression.Streams; using K4os.Hash.Crc; using Wabbajack.Common; +using Wabbajack.Common.FileSignatures; +using Wabbajack.ImageHashing; namespace Wabbajack.VirtualFileSystem { @@ -47,6 +49,7 @@ namespace Wabbajack.VirtualFileSystem public FullPath FullPath { get; private set; } public Hash Hash { get; internal set; } + public PHash PerceptualHash { get; internal set; } public ExtendedHashes ExtendedHashes { get; set; } public long Size { get; internal set; } @@ -177,16 +180,17 @@ namespace Wabbajack.VirtualFileSystem private IndexedVirtualFile ToIndexedVirtualFile() { - return new IndexedVirtualFile + return new() { Hash = Hash, + PerceptualHash = PerceptualHash, Name = Name, Children = Children.Select(c => c.ToIndexedVirtualFile()).ToList(), Size = Size }; } - - + + public static async Task Analyze(Context context, VirtualFile parent, IStreamFactory extractedFile, IPath relPath, int depth = 0) { @@ -219,9 +223,12 @@ namespace Wabbajack.VirtualFileSystem Size = stream.Length, LastModified = extractedFile.LastModifiedUtc.AsUnixTime(), LastAnalyzed = DateTime.Now.AsUnixTime(), - Hash = hash + Hash = hash, }; + if (Consts.TextureExtensions.Contains(relPath.FileName.Extension)) + self.PerceptualHash = await PHash.FromStream(stream, relPath.FileName.Extension, false); + self.FillFullPath(depth); if (context.UseExtendedHashes) diff --git a/Wabbajack.VirtualFileSystem/Wabbajack.VirtualFileSystem.csproj b/Wabbajack.VirtualFileSystem/Wabbajack.VirtualFileSystem.csproj index 362aa1b2..39d81a7d 100644 --- a/Wabbajack.VirtualFileSystem/Wabbajack.VirtualFileSystem.csproj +++ b/Wabbajack.VirtualFileSystem/Wabbajack.VirtualFileSystem.csproj @@ -12,6 +12,7 @@ +