Mirrored_Repos/wabbajack
1
0
Fork 0
mirror of https://github.com/wabbajack-tools/wabbajack.git synced 2024-08-30 18:42:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP, hooking it into the VFS cache

Browse Source
This commit is contained in:
Timothy Baldridge 2021-06-15 23:16:25 -06:00
parent c9a2ea8180
commit abd142623f
13 changed files with 286 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Wabbajack.Common/Consts.cs
View File

@ -172,5 +172,6 @@ namespace Wabbajack.Common
public static long UPLOADED_FILE_BLOCK_SIZE = (long)1024 * 1024 * 2;
public static HashSet<Extension> TextureExtensions = new() {new Extension(".dds"), new Extension(".tga")};
}
}

47
Wabbajack.ImageHashing.Test/ImageLoadingTests.cs
View File

@ -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));
}
}
}

BIN
Wabbajack.ImageHashing.Test/Resources/test-dxt5-recompressed.dds Normal file
View File

Binary file not shown.

BIN
Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7-vflip.dds Normal file
View File

Binary file not shown.

BIN
Wabbajack.ImageHashing.Test/Resources/test-dxt5-small-bc7.dds Normal file
View File

Binary file not shown.

11
Wabbajack.ImageHashing.Test/Wabbajack.ImageHashing.Test.csproj
View File

@ -4,6 +4,8 @@
<TargetFramework>net5.0-windows</TargetFramework>
<IsPackable>false</IsPackable>
<Nullable>enable</Nullable>
</PropertyGroup>
<ItemGroup>
@ -27,6 +29,15 @@
<None Update="Resources\test-dxt5.dds">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</None>
<None Update="Resources\test-dxt5-small-bc7.dds">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</None>
<None Update="Resources\test-dxt5-small-bc7-vflip.dds">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</None>
<None Update="Resources\test-dxt5-recompressed.dds">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</None>
</ItemGroup>
</Project>

125
Wabbajack.ImageHashing/DDSImage.cs
View File

@ -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<DXGI_FORMAT> CompressedTypes = new HashSet<DXGI_FORMAT>()
{
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();
}
}
}
}

18
Wabbajack.ImageHashing/IImage.cs
View File

@ -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
}
}

89
Wabbajack.ImageHashing/PHash.cs Normal file
View File

@ -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<PHash> 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();
}
}
}

4
Wabbajack.ImageHashing/Wabbajack.ImageHashing.csproj
View File

@ -2,11 +2,13 @@
<PropertyGroup>
<TargetFramework>net5.0-windows</TargetFramework>
<AllowUnsafeBlocks>true</AllowUnsafeBlocks>
<Nullable>enable</Nullable>
</PropertyGroup>
<ItemGroup>
<PackageReference Include="DirectXTexNet" Version="1.0.2" />
<PackageReference Include="Magick.NET.Core" Version="7.0.1" />
<PackageReference Include="Shipwreck.Phash" Version="0.5.0" />
</ItemGroup>
<ItemGroup>

7
Wabbajack.VirtualFileSystem/IndexedVirtualFile.cs
View File

@ -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<IndexedVirtualFile> Children { get; set; } = new List<IndexedVirtualFile>();
public List<IndexedVirtualFile> 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<IndexedVirtualFile>();

15
Wabbajack.VirtualFileSystem/VirtualFile.cs
View File

@ -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<VirtualFile> 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)

1
Wabbajack.VirtualFileSystem/Wabbajack.VirtualFileSystem.csproj
View File

@ -12,6 +12,7 @@
<ItemGroup>
<ProjectReference Include="..\Compression.BSA\Compression.BSA.csproj" />
<ProjectReference Include="..\Wabbajack.Common\Wabbajack.Common.csproj" />
<ProjectReference Include="..\Wabbajack.ImageHashing\Wabbajack.ImageHashing.csproj" />
</ItemGroup>
<ItemGroup>
<PackageReference Include="Genbox.AlphaFS" Version="2.2.2.1" />