This commit is contained in:
Timothy Baldridge 2021-06-16 15:07:16 -06:00
parent abd142623f
commit 6706d5bfb6
4 changed files with 91 additions and 16 deletions

View File

@ -12,10 +12,11 @@ namespace Wabbajack.ImageHashing
{
public class DDSImage
{
private DDSImage(ScratchImage img, TexMetadata metadata)
private DDSImage(ScratchImage img, TexMetadata metadata, Extension ext)
{
_image = img;
_metaData = metadata;
_extension = ext;
}
private static Extension DDSExtension = new(".dds");
@ -29,7 +30,7 @@ namespace Wabbajack.ImageHashing
var img = TexHelper.Instance.LoadFromDDSFile(file.ToString(), DDS_FLAGS.NONE);
return new DDSImage(img, img.GetMetadata());
return new DDSImage(img, img.GetMetadata(), new Extension(".dds"));
}
public static DDSImage FromDDSMemory(byte[] data)
@ -39,7 +40,7 @@ namespace Wabbajack.ImageHashing
fixed (byte* ptr = data)
{
var img = TexHelper.Instance.LoadFromDDSMemory((IntPtr)ptr, data.Length, DDS_FLAGS.NONE);
return new DDSImage(img, img.GetMetadata());
return new DDSImage(img, img.GetMetadata(), new Extension(".dds"));
}
}
}
@ -51,7 +52,7 @@ namespace Wabbajack.ImageHashing
fixed (byte* ptr = data)
{
var img = TexHelper.Instance.LoadFromTGAMemory((IntPtr)ptr, data.Length);
return new DDSImage(img, img.GetMetadata());
return new DDSImage(img, img.GetMetadata(), new Extension(".tga"));
}
}
}
@ -95,11 +96,24 @@ namespace Wabbajack.ImageHashing
DXGI_FORMAT.BC7_UNORM_SRGB,
};
private Extension _extension;
public ImageState ImageState()
{
return new()
{
Width = _metaData.Width,
Height = _metaData.Height,
PerceptualHash = PerceptionHash()
};
}
public PHash PerceptionHash()
{
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);
@ -109,8 +123,7 @@ namespace Wabbajack.ImageHashing
{
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)
@ -121,7 +134,6 @@ namespace Wabbajack.ImageHashing
if (widthV != 512) return;
var y = line.ToInt32();
data.Add((widthV, y));
for (int i = 0; i < widthV; i++)
{

View File

@ -0,0 +1,33 @@
using System.IO;
using DirectXTexNet;
using Wabbajack.Common;
namespace Wabbajack.ImageHashing
{
public class ImageState
{
public int Width { get; set; }
public int Height { get; set; }
public DXGI_FORMAT Format { get; set; }
public PHash PerceptualHash { get; set; }
public static ImageState Read(BinaryReader br)
{
return new()
{
Width = br.ReadUInt16(),
Height = br.ReadUInt16(),
Format = (DXGI_FORMAT)br.ReadByte(),
PerceptualHash = PHash.Read(br)
};
}
public void Write(BinaryWriter bw)
{
bw.Write((ushort)Width);
bw.Write((ushort)Height);
bw.Write((byte)Format);
PerceptualHash.Write(bw);
}
}
}

View File

@ -11,12 +11,24 @@ namespace Wabbajack.ImageHashing
{
private const int SIZE = 40;
private readonly byte[] _data;
private readonly int _hash;
private PHash(byte[]? data)
private PHash(byte[] data)
{
_data = data ?? new byte[SIZE];
_data = data;
if (_data.Length != SIZE)
throw new DataException();
long h = 0;
h |= _data[0];
h <<= 8;
h |= _data[1];
h <<= 8;
h |= _data[2];
h <<= 8;
h |= _data[3];
h <<= 8;
_hash = (int)h;
}
public static PHash FromBase64(string base64)
@ -34,7 +46,10 @@ namespace Wabbajack.ImageHashing
public void Write(BinaryWriter br)
{
br.Write((_data?.Length ?? 0) == 0 ? new byte[SIZE] : _data);
if (_hash == 0)
br.Write(new byte[SIZE]);
else
br.Write(_data);
}
public static PHash FromDigest(Digest digest)

View File

@ -1,5 +1,6 @@
using System.Collections.Generic;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
@ -18,7 +19,7 @@ namespace Wabbajack.VirtualFileSystem
public IPath Name { get; set; }
public Hash Hash { get; set; }
public PHash PerceptualHash { get; set; }
public ImageState? ImageState { get; set; }
public long Size { get; set; }
public List<IndexedVirtualFile> Children { get; set; } = new();
@ -26,7 +27,15 @@ namespace Wabbajack.VirtualFileSystem
{
bw.Write(Name.ToString());
bw.Write((ulong)Hash);
PerceptualHash.Write(bw);
if (ImageState == null)
bw.Write(false);
else
{
bw.Write(true);
ImageState.Write(bw);
}
bw.Write(Size);
bw.Write(Children.Count);
foreach (var file in Children)
@ -35,7 +44,8 @@ namespace Wabbajack.VirtualFileSystem
public void Write(Stream s)
{
using var bw = new BinaryWriter(s, Encoding.UTF8, true);
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)
@ -48,9 +58,13 @@ namespace Wabbajack.VirtualFileSystem
{
Name = (RelativePath)br.ReadString(),
Hash = Hash.FromULong(br.ReadUInt64()),
PerceptualHash = PHash.Read(br),
Size = br.ReadInt64(),
};
if (br.ReadBoolean())
ivf.ImageState = ImageState.Read(br);
ivf.Size = br.ReadInt64();
var lst = new List<IndexedVirtualFile>();
ivf.Children = lst;
var count = br.ReadInt32();
@ -64,7 +78,8 @@ namespace Wabbajack.VirtualFileSystem
public static IndexedVirtualFile Read(Stream s)
{
using var br = new BinaryReader(s);
using var cs = new GZipStream(s, CompressionMode.Decompress, true);
using var br = new BinaryReader(cs);
var ivf = new IndexedVirtualFile
{
Size = br.ReadInt64(),