KSP-Toolkit/Scripts/dds.py

import os
import subprocess
import struct
import re
import math
import enum
import s3tc

if os.name == 'nt':
    convertcmd = "convert.exe"
    compresscmd = "nvcompress.exe"
    decompresscmd = "nvdecompress.exe"
    infocmd = "nvddsinfo.exe"
else:
    convertcmd = "convert"
    compresscmd = "nvcompress"
    decompresscmd = "nvdecompress"
    infocmd = "nvddsinfo"

def alpha(file):
    result = subprocess.run([convertcmd, file, "-resize", "1x1",
                             "-format", "%[fx:int(255*a+.5)]", "info:-"], capture_output=True)
    result.check_returncode()
    return int(result.stdout)


def flip(file, output):
    result = subprocess.run(
        [convertcmd, '-flip', file, output], capture_output=True)
    result.check_returncode()


def nvcompress(format, file, output, mips=True):
    args = [format]
    if not mips:
        args.append('-nomips')

    result = subprocess.run([compresscmd] + args +
                            [file, output], capture_output=True)
    result.check_returncode()


def nvdecompress(file, output):
    result = subprocess.run([decompresscmd, file, output], capture_output=True)
    result.check_returncode()


def nvinfo(file):
    result = subprocess.run([infocmd, file], capture_output=True)
    result.check_returncode()

    info = {
        "format": re.search(r"FourCC: '(.{4})'", str(result.stdout)).group(1)
    }

    return info

class DDSFlags(enum.IntFlag):
    CAPS = 0x1
    HEIGHT = 0x2
    WIDTH = 0x4
    PITCH = 0x8
    PIXEL_FORMAT = 0x1000
    MIPMAPCOUNT = 0x20000
    LINEAR_SIZE = 0x80000
    DEPTH = 0x800000

class DDSPixelFlags(enum.IntFlag):
    ALPHAPIXELS = 0x1
    ALPHA = 0x2
    FOURCC = 0x4
    RGB = 0x40
    YUV = 0x200
    LUMINANCE = 0x20000


texture_formats = {
    'DXT1' : s3tc.DXT1Texture,
    'DXT5' : s3tc.DXT5Texture
}

class DDSFile:
    def __init__(self, file = None):
        if (file):
            self.read(file)

    def read(self, file):
        self.read_header(file)
        self.read_data(file)

    def write(self, file):
        self.write_header(file)
        self.write_data(file)

    def read_header(self, file):
        # read magic bytes
        if file.read(4).decode() != 'DDS ':
            raise ValueError("File is not a valid DDS file: Incorrect magic bytes")

        # read header
        header = struct.unpack('<7I44x', file.read(72))

        # check constant
        if header[0] != 124:
            raise ValueError("File is not a valid DDS file: Incorrect header size")

        # read basic info
        self.flags = DDSFlags(header[1])
        self.height, self.width, self.linear_size, self.depth, self.mipmapcount = header[2:7]

        # read pixelformat
        pixelformat = struct.unpack('<2I4s5I', file.read(32))

        # check constant
        if pixelformat[0] != 32:
            raise ValueError("File is not a valid DDS file: Incorrect pixelformat size")

        self.pixel_flags = DDSPixelFlags(pixelformat[1])
        self.four_cc = pixelformat[2].decode()
        self.rgb_bit_count, self.r_bitmask, self.g_bitmask, self.b_bitmask, self.a_bitmask = pixelformat[3:8]

        # read caps
        self.dw_caps, self.dw_caps2, self.dw_caps3, self.dw_caps4 = struct.unpack('<4I4x', file.read(20))

        # read DX10 header
        self.is_DX10 = (self.four_cc == 'DX10')
        if self.is_DX10:
            self.read_dx10header(file)

    def read_dx10header(self, file):
        dx10header = file.read(20)
        #TODO: actually do something with the DX10 data

    def write_header(self, file):
        file.write(bytes('DDS ', 'utf-8'))

        # write header
        file.write(struct.pack(
            '<7I44x', 124, int(self.flags), self.height, self.width, 
            self.linear_size, self.depth, self.mipmapcount))

        # write pixelformat
        file.write(struct.pack(
            '<2I4s5I', 32, int(self.pixel_flags), bytes(self.four_cc, 'utf-8'), self.rgb_bit_count,
            self.r_bitmask, self.g_bitmask, self.b_bitmask, self.a_bitmask))

        # write caps
        file.write(struct.pack(
            '<4I4x', self.dw_caps, self.dw_caps2, self.dw_caps3, self.dw_caps4))

        #TODO: handle DX10 data
        if self.is_DX10:
            file.write(struct.pack('20x'))

    def read_data(self, file):
        self.textures = []

        if DDSFlags.MIPMAPCOUNT in self.flags:
            mips = self.mipmapcount
        else:
            mips = 1

        if DDSPixelFlags.FOURCC in self.pixel_flags:
            textureType = texture_formats[self.four_cc]
            for mip in range(mips):
                width = math.ceil(self.width / (2 ** mip))
                height = math.ceil(self.height / (2 ** mip))
                self.textures.append(textureType(width, height, file))

    def write_data(self, file):
        for texture in self.textures:
            texture.write(file)