From 8b6ea69300b438d1f59fcaa2c17b44c9a1e1c45c Mon Sep 17 00:00:00 2001 From: drewcassidy Date: Thu, 8 Apr 2021 21:26:58 -0700 Subject: [PATCH] DDS file class --- quicktex/dds.py | 441 +++++++++++++++++----------------------- quicktex/image_utils.py | 31 +-- 2 files changed, 186 insertions(+), 286 deletions(-) diff --git a/quicktex/dds.py b/quicktex/dds.py index f4be690..bb58041 100644 --- a/quicktex/dds.py +++ b/quicktex/dds.py @@ -1,218 +1,90 @@ from __future__ import annotations import enum +import os import struct import typing +import quicktex.image_utils +import quicktex.s3tc.bc1 +import quicktex.s3tc.bc3 +import quicktex.s3tc.bc4 +import quicktex.s3tc.bc5 +from PIL import Image -@typing.final -class PixelFormat: - """ - DDS header surface format. - - For more information, see microsoft documentation for `DDS_PIXELFORMAT `_. - """ - size = 32 - """The size of a PixelFormat block in bytes.""" - - class Flags(enum.IntFlag): - """Values which indicate what type of data is in the surface.""" - - ALPHAPIXELS = 0x1 - """Texture contains alpha data (:py:attr:`~PixelFormat.pixel_bitmasks[3]` contains valid data).""" - - ALPHA = 0x2 - """Used in some older DDS files for alpha channel only uncompressed data - (:py:attr:`~PixelFormat.pixel_size` contains the alpha channel bitcount; :py:attr:`~PixelFormat.pixel_bitmasks[3]` contains valid data).""" - - FOURCC = 0x4 - """Texture contains compressed RGB data; :py:attr:`~PixelFormat.four_cc` contains valid data.""" - - RGB = 0x40 - """Texture contains uncompressed RGB data; :py:attr:`~PixelFormat.pixel_size` and the RGB masks - (:py:attr:`~PixelFormat.pixel_bitmasks[0:3]`) contain valid data.""" - - YUV = 0x200 - """Used in some older DDS files for YUV uncompressed data - (:py:attr:`~PixelFormat.pixel_size` contains the YUV bit count; :py:attr:`~PixelFormat.pixel_bitmasks[0]` contains the Y mask, - :py:attr:`~PixelFormat.pixel_bitmasks[1]` contains the U mask, :py:attr:`~PixelFormat.pixel_bitmasks[2]` contains the V mask).""" - - LUMINANCE = 0x20000 - """Used in some older DDS files for single channel color uncompressed data (:py:attr:`~PixelFormat.pixel_size` - contains the luminance channel bit count; :py:attr:`~PixelFormat.pixel_bitmasks[0]` contains the channel mask). - Can be combined with :py:attr:`ALPHAPIXELS` for a two channel uncompressed DDS file.""" - - def __init__(self): - self.flags: PixelFormat.Flags = PixelFormat.Flags.FOURCC - """Flags representing which data is valid.""" - - self.four_cc: str = "NONE" - """FourCC code of the texture format. Valid texture format strings are ``DXT1``, ``DXT2``, ``DXT3``, ``DXT4``, or ``DXT5``. - If a DirectX 10 header is used, this is ``DX10``.""" - - self.pixel_size: int = 0 - """Number of bits in each pixel if the texture is uncompressed""" - - self.pixel_bitmasks: typing.Tuple[int, int, int, int] = (0, 0, 0, 0) - """Tuple of bitmasks for each channel""" - - @staticmethod - def from_bytes(data) -> PixelFormat: - """ - Create a new PixelFormat object from a bytes-like object - - :param data: A bytes-like object holding the raw data to unpack - :return: An unpacked PixelFormat object - """ - assert len(data) == PixelFormat.size, "Incorrect number of bytes in input." - - unpacked = struct.unpack('<2I4s5I', data) - assert unpacked[0] == PixelFormat.size, "Incorrect pixelformat size." - - pf = PixelFormat() - pf.flags = PixelFormat.Flags(unpacked[1]) - pf.four_cc = unpacked[2].decode() - pf.pixel_size = unpacked[3] - pf.pixel_bitmasks = unpacked[4:8] - return pf - - @staticmethod - def from_file(file: typing.BinaryIO) -> PixelFormat: - """ - Create a new PixelFormat object from a file. The file position will be advanced by 32 bytes. - - :param file: A file-like object to read from. - :return: An unpacked PixelFormat object - """ - assert file.readable(), "Input file is not readable." - data = file.read(PixelFormat.size) - return PixelFormat.from_bytes(data) - - def to_bytes(self) -> bytes: - """ - Write the PixelFormat object to a bytes object. - - :return: The packed PixelFormat object - """ - data = struct.pack('<2I4s5I', 32, int(self.flags), bytes(self.four_cc, 'ascii'), self.pixel_size, *self.pixel_bitmasks) - assert len(data) == PixelFormat.size - return data +class DDSFormat: + def __init__(self, name: str, texture, encoder, decoder, four_cc: str = None): + self.four_cc = four_cc + self.decoder = decoder + self.encoder = encoder + self.texture = texture + self.name = name -@typing.final -class DDSHeader: - """ - Header for a microsoft DDS file +dds_formats = [ + DDSFormat('BC1', quicktex.s3tc.bc1.BC1Texture, quicktex.s3tc.bc1.BC1Encoder, quicktex.s3tc.bc1.BC1Decoder, 'DXT1'), + DDSFormat('BC3', quicktex.s3tc.bc3.BC3Texture, quicktex.s3tc.bc3.BC3Encoder, quicktex.s3tc.bc3.BC3Decoder, 'DXT5'), + DDSFormat('BC4', quicktex.s3tc.bc4.BC4Texture, quicktex.s3tc.bc4.BC4Encoder, quicktex.s3tc.bc4.BC4Decoder, 'ATI1'), + DDSFormat('BC5', quicktex.s3tc.bc5.BC5Texture, quicktex.s3tc.bc5.BC5Encoder, quicktex.s3tc.bc5.BC5Decoder, 'ATI2'), +] - For more information, see microsoft documentation for `DDS_HEADER `_. - """ - size = 124 - """The size of a DDS header in bytes.""" - class Flags(enum.IntFlag): - """Flags to indicate which members contain valid data.""" +class PFFlags(enum.IntFlag): + """Values which indicate what type of data is in the surface.""" - CAPS = 0x1 - """Required in every .dds file.""" + ALPHAPIXELS = 0x1 + """Texture contains alpha data (:py:attr:`~PixelFormat.pixel_bitmasks[3]` contains valid data).""" - HEIGHT = 0x2 - """Required in every .dds file.""" + ALPHA = 0x2 + """Used in some older DDS files for alpha channel only uncompressed data + (:py:attr:`~PixelFormat.pixel_size` contains the alpha channel bitcount; :py:attr:`~PixelFormat.pixel_bitmasks[3]` contains valid data).""" - WIDTH = 0x4 - """Required in every .dds file.""" + FOURCC = 0x4 + """Texture contains compressed RGB data; :py:attr:`~PixelFormat.four_cc` contains valid data.""" - PITCH = 0x8 - """Required when :py:attr:`~DDSHeader.pitch` is provided for an uncompressed texture.""" + RGB = 0x40 + """Texture contains uncompressed RGB data; :py:attr:`~PixelFormat.pixel_size` and the RGB masks + (:py:attr:`~PixelFormat.pixel_bitmasks[0:3]`) contain valid data.""" - PIXEL_FORMAT = 0x1000 - """Required in every .dds file.""" + YUV = 0x200 + """Used in some older DDS files for YUV uncompressed data + (:py:attr:`~PixelFormat.pixel_size` contains the YUV bit count; :py:attr:`~PixelFormat.pixel_bitmasks[0]` contains the Y mask, + :py:attr:`~PixelFormat.pixel_bitmasks[1]` contains the U mask, :py:attr:`~PixelFormat.pixel_bitmasks[2]` contains the V mask).""" - MIPMAPCOUNT = 0x20000 - """Required when :py:attr:`~DDSHeader.mipmap_count` is provided for a mipmapped texture.""" + LUMINANCE = 0x20000 + """Used in some older DDS files for single channel color uncompressed data (:py:attr:`~PixelFormat.pixel_size` + contains the luminance channel bit count; :py:attr:`~PixelFormat.pixel_bitmasks[0]` contains the channel mask). + Can be combined with :py:attr:`ALPHAPIXELS` for a two channel uncompressed DDS file.""" - LINEAR_SIZE = 0x80000 - """Required when :py:attr:`~DDSHeader.pitch` is provided for a compressed texture.""" - DEPTH = 0x800000 - """Required when :py:attr:`~DDSHeader.depth` is provided for a depth texture.""" +class DDSFlags(enum.IntFlag): + """Flags to indicate which members contain valid data.""" - REQUIRED = CAPS | HEIGHT | WIDTH | PIXEL_FORMAT + CAPS = 0x1 + """Required in every .dds file.""" - def __init__(self): - self.flags: DDSHeader.Flags = DDSHeader.Flags.REQUIRED - """Flags to indicate which members contain valid data.""" + HEIGHT = 0x2 + """Required in every .dds file.""" - self.dimensions: typing.Tuple[int, int] = (0, 0) - """Width and height of the texture or its first mipmap""" + WIDTH = 0x4 + """Required in every .dds file.""" - self.pitch: int = 0 - """The pitch or number of bytes per scan line in an uncompressed texture; - the total number of bytes in the top level texture for a compressed texture.""" + PITCH = 0x8 + """Required when :py:attr:`~DDSHeader.pitch` is provided for an uncompressed texture.""" - self.depth: int = 0 - """Depth of a volume texture (in pixels), otherwise unused.""" + PIXEL_FORMAT = 0x1000 + """Required in every .dds file.""" - self.mipmap_count: int = 0 - """Number of mipmap levels, otherwise unused.""" + MIPMAPCOUNT = 0x20000 + """Required when :py:attr:`~DDSHeader.mipmap_count` is provided for a mipmapped texture.""" - self.pixel_format: PixelFormat = PixelFormat() - """The pixel format""" + LINEAR_SIZE = 0x80000 + """Required when :py:attr:`~DDSHeader.pitch` is provided for a compressed texture.""" - self.caps: typing.Tuple[int, int, int, int] = (0, 0, 0, 0) - """Specifies the complexity of the surfaces stored.""" + DEPTH = 0x800000 + """Required when :py:attr:`~DDSHeader.depth` is provided for a depth texture.""" - @staticmethod - def from_bytes(data) -> DDSHeader: - """ - Create a new DDS Header from a bytes-like object - - :param data: A bytes-like object holding the raw data to unpack - :return: An unpacked DDS header - """ - assert len(data) == PixelFormat.size, "Incorrect number of bytes in input." - - unpacked = struct.unpack('<7I44x', data[0:72]) - assert unpacked[0] == DDSHeader.size, "Incorrect DDS header size." - - header = DDSHeader() - header.flags = DDSHeader.Flags(unpacked[1]) - header.dimensions = unpacked[2:4:-1] - header.pitch, header.depth, header.mipmapcount = unpacked[4:7] - - header.pixel_format = PixelFormat.from_bytes(data[72:104]) - - header.caps = struct.unpack('<4I4x', data[104:124]) - - assert all([val > 0 for val in header.dimensions]), "Image size is zero" - - return header - - @staticmethod - def from_file(file: typing.BinaryIO) -> DDSHeader: - """ - Create a new DDS header from a file. the file position will be advanced by 124 bytes. - - :param file: A file-like object to read from. - :return: An unpacked DDS header - """ - assert file.readable(), "Input file is not readable." - return DDSHeader.from_bytes(file.read(DDSHeader.size)) - - def to_bytes(self) -> bytes: - """ - Write the DDS header to a bytes object. - - :return: The packed DDS header - """ - data = b'' - # write header - data += struct.pack('<7I44x', DDSHeader.size, int(self.flags), self.dimensions[1], self.dimensions[0], - self.pitch, self.depth, self.mipmap_count) - data += self.pixel_format.to_bytes() - data += struct.pack('<4I4x', *self.caps) - - assert len(data) == DDSHeader.size - return data + TEXTURE = CAPS | HEIGHT | WIDTH | PIXEL_FORMAT @typing.final @@ -229,87 +101,142 @@ class DDSFile: extension = 'dds' """Extension for a DDS file.""" - def __init__(self): - self.header: DDSHeader = DDSHeader() - """The DDS file's header object""" + header_bytes = 124 + """The size of a DDS header in bytes.""" - self.textures: typing.List[bytes] = [] + def __init__(self): + self.flags: DDSFlags = DDSFlags.TEXTURE + """Flags to indicate which members contain valid data.""" + + self.size: typing.Tuple[int, int] = (0, 0) + """Width and height of the texture or its first mipmap""" + + self.pitch: int = 0 + """The pitch or number of bytes per row in an uncompressed texture; + the total number of bytes in the top level texture for a compressed texture.""" + + self.depth: int = 1 + """Depth of a volume texture (in pixels), otherwise unused.""" + + self.mipmap_count: int = 1 + """Number of mipmap levels, otherwise unused.""" + + self.pf_flags: PFFlags = PFFlags.FOURCC + """Flags representing which pixel format data is valid.""" + + self.four_cc: str = "NONE" + """FourCC code of the texture format. Valid texture format strings are ``DXT1``, ``DXT2``, ``DXT3``, ``DXT4``, or ``DXT5``. + If a DirectX 10 header is used, this is ``DX10``.""" + + self.pixel_size: int = 0 + """Number of bits in each pixel if the texture is uncompressed""" + + self.pixel_bitmasks: typing.Tuple[int, int, int, int] = (0, 0, 0, 0) + """Tuple of bitmasks for each channel""" + + self.caps: typing.Tuple[int, int, int, int] = (0, 0, 0, 0) + """Specifies the complexity of the surfaces stored.""" + + self.textures: typing.List = [] """A list of bytes objects for each texture in the file""" - @staticmethod - def from_file(file: typing.BinaryIO) -> DDSFile: - """ - Create a new DDSFile object from the contents of a file. + self.format: DDSFormat = DDSFormat('NONE', None, None, None) + """The format used by this dds file""" - :param file: A file-like object to read from. - :return: An unpacked DDS file + def save(self, path: os.PathLike) -> None: """ - assert file.readable(), "Input file is not readable." + Save the DDSFile to a file + :param path: string or path-like object to write to + """ + with open(path, 'wb') as file: + self.size = self.textures[0].size + self.pitch = self.textures[0].nbytes + self.mipmap_count = len(self.textures) + + assert quicktex.image_utils.mip_sizes(self.size, self.mipmap_count) == [tex.size for tex in self.textures], 'incorrect mipmap sizes' + + file.write(DDSFile.magic) + + # WRITE HEADER + file.write(struct.pack('<7I44x', DDSFile.header_bytes, int(self.flags), self.size[1], self.size[0], self.pitch, self.depth, self.mipmap_count)) + file.write(struct.pack('<2I4s5I', 32, int(self.flags), bytes(self.four_cc, 'ascii'), self.pixel_size, *self.pixel_bitmasks)) + file.write(struct.pack('<4I4x', *self.caps)) + + assert file.tell() == 4 + DDSFile.header_bytes, 'error writing file: incorrect header size' + + for texture in self.textures: + file.write(texture) + + def decode(self, mip: int = 0, *args, **kwargs) -> Image.Image: + """ + Decode a single texture in the file to images + :param mip: the mip level to decode. Default: 0 + :return: The decoded image + """ + decoder = self.format.decoder(*args, **kwargs) + texture = decoder.decode(self.textures[mip]) + return Image.frombuffer('RGBA', texture.size, texture) + + def decode_all(self, *args, **kwargs) -> typing.List[Image.Image]: + """ + Decade all textures in the file to images + :return: the decoded images + """ + decoder = self.format.decoder(*args, **kwargs) + textures = [decoder.decode(encoded) for encoded in self.textures] + return [Image.frombuffer('RGBA', tex.size, tex) for tex in textures] + + +def read(path: os.PathLike) -> DDSFile: + with open(path, 'rb') as file: assert file.read(4) == DDSFile.magic, "Incorrect magic bytes in DDS file." dds = DDSFile() - dds.header = DDSHeader.from_file(file) - four_cc = dds.header.pixel_format.four_cc - assert four_cc == 'DXT1' or four_cc == 'DXT5' - block_size = 8 if four_cc == 'DXT1' else 16 - mip_count = dds.header.mipmap_count if DDSHeader.Flags.MIPMAPCOUNT in dds.header.flags else 1 + # READ HEADER + assert struct.unpack(' int: - """ - Get the number of mipmaps in the dds file. - - :return: The number of mipmaps - """ - if DDSHeader.Flags.MIPMAPCOUNT in self.header.flags: - mips = self.header.mipmap_count - assert mips > 0 - return mips - else: - return 1 - - def image_dimensions(self, mip: int = 0) -> typing.Tuple[int, int]: - """ - Calculate the dimensions of a mip level in pixels. - - :param int mip: which mip level to calculate the dimensions of, between 0 and :py:attr:`~DDSHeader.mipmap_count` exclusive. - :return: the dimensions of the selected mipmap in pixels - """ - assert self.mipmap_count() > mip >= 0, "Invalid mip level" - return tuple([max(size // (2 ** mip), 1) for size in self.header.dimensions]) - - def image_block_dimensions(self, mip: int = 0) -> typing.Tuple[int, int]: - """ - Calculate the dimensions of a mip level in 4x4 blocks. Only relevent for compressed textures. - - :param int mip: Whick mip level to calculate the dimensions of, between 0 and :py:attr:`~DDSHeader.mipmap_count` exclusive. - :return: the dimensions of the selected mipmap in blocks - """ - dimensions = self.image_dimensions(mip) - - assert all(size > 0 for size in dimensions) - return tuple([max((size + 3) // 4, 1) for size in dimensions]) diff --git a/quicktex/image_utils.py b/quicktex/image_utils.py index 8b4919d..2c18f0c 100644 --- a/quicktex/image_utils.py +++ b/quicktex/image_utils.py @@ -5,35 +5,6 @@ import typing import math -def pad(source: Image.Image, block_dimensions=(4, 4)) -> Image.Image: - """ - Pad an image to be divisible by a specific block size. The input image is repeated into the unused areas so that bilinar filtering works correctly. - - :param source: Input image to add padding to. This will not be modified. - :param block_dimensions: The size of a single block that the output must be divisible by. - :return: A new image with the specified padding added. - """ - - assert all([dim > 0 for dim in block_dimensions]), "Invalid block size" - - padded_dimensions = tuple([ - math.ceil(i_dim / b_dim) * b_dim - for i_dim in source.size - for b_dim in block_dimensions - ]) - - if padded_dimensions == source.size: - # no padding is necessary - return source - - output = Image.new(source.mode, padded_dimensions) - for x in range(math.ceil(padded_dimensions[0] / source.width)): - for y in range(math.ceil(padded_dimensions[1] / source.height)): - output.paste(source, (x * source.width, y * source.height)) - - return output - - def mip_sizes(dimensions: typing.Tuple[int, int], mip_count: typing.Optional[int] = None) -> typing.List[typing.Tuple[int, int]]: """ Create a chain of mipmap sizes for a given source source size, where each source is half the size of the one before. @@ -62,3 +33,5 @@ def mip_sizes(dimensions: typing.Tuple[int, int], mip_count: typing.Optional[int break # we've reached a 1x1 mip and can get no smaller return chain + +