Les algorithmes S3TC : la compression de textures

from typing import List, Tuple
from io import BytesIO

def couleur_16bits_vers_rgba(entree : int) -> Tuple[int, int, int, int]:
    
    return (
        int(( entree >> (6 + 5) ) / 0b11111 * 255),
        int(( (entree >> 5) & 0b111111 ) / 0b111111 * 255),
        int(( entree & 0b11111 ) / 0b11111 * 255),
        255
    )


"""
    Notre décodeur se trouve ici. Les dimensions de l'image doivent être
    définies explicitement, puisque autrement on ne sait pas combien de
    blocs sont horizontaux et combien sont verticaux.
"""

def decoder_dxt1(hauteur_px : int, largeur_px : int, entree : bytes) -> List[List[Tuple[int, int, int, int]]]:
    
    lecteur = BytesIO(entree)
    
    sortie : List[List[Tuple[int, int, int, int]]] = []
    
    for base_position_y in range(0, hauteur_px, 4):
        
        for base_position_x in range(0, largeur_px, 4):
    
            couleur_1_brute = int.from_bytes(lecteur.read(2), 'little')
            couleur_2_brute = int.from_bytes(lecteur.read(2), 'little')
            
            couleur_1 = couleur_16bits_vers_rgba(couleur_1_brute)
            couleur_2 = couleur_16bits_vers_rgba(couleur_2_brute)
            
            if couleur_1_brute > couleur_2_brute:
                couleur_3 = tuple(
                    int(2/3 * couleur_1[composante] + 1/3 * couleur_2[composante])
                    for composante in range(4)
                )
                couleur_4 = tuple(
                    int(1/3 * couleur_1[composante] + 2/3 * couleur_2[composante])
                    for composante in range(4)
                )
            else:
                couleur_3 = tuple(
                    int(1/2 * couleur_1[composante] + 1/2 * couleur_2[composante])
                    for composante in range(4)
                )
                couleur_4 = (0, 0, 0, 0) # Noir transparent
                
            
            couleurs = [couleur_1, couleur_2, couleur_3, couleur_4]
            
            for position_y in range(base_position_y, base_position_y + 4):                
                
                ligne_encodee = lecteur.read(1)[0]
                
                while len(sortie) <= position_y:
                    sortie.append([]) # Ajouter une ligne de pixels à notre sortie
                
                sortie[position_y] += [
                    couleurs[ligne_encodee        & 0b11],
                    couleurs[(ligne_encodee >> 2) & 0b11],
                    couleurs[(ligne_encodee >> 4) & 0b11],
                    couleurs[ligne_encodee >> 6],
                ]
    
    return sortie
        

# D'après https://docs.microsoft.com/en-us/windows/win32/direct3ddds/dds-pixelformat

class DDS_PIXELFORMAT(LittleEndianStructure):
    
    _fields_ = [
        ('dwSize', c_uint32),
        ('dwFlags', c_uint32),
        ('dwFourCC', c_char * 4),
        ('dwRGBBitCount', c_uint32),
        ('dwRBitMask', c_uint32),
        ('dwGBitMask', c_uint32),
        ('dwBBitMask', c_uint32),
        ('dwABitMask', c_uint32)
    ]
    
    _pack_ = True

# D'après https://docs.microsoft.com/en-us/windows/win32/direct3ddds/dds-header

class DDS_HEADER(LittleEndianStructure):

    _fields_ = [
        ('dwMagic', c_char * 4),
        ('dwSize', c_uint32),
        ('dwFlags', c_uint32),
        ('dwHeight', c_uint32),
        ('dwWidth', c_uint32),
        ('dwPitchOrLinearSize', c_uint32),
        ('dwDepth', c_uint32),
        ('dwMipMapCount', c_uint32),
        ('dwReserved', c_uint32 * 11),
        ('ddspf', DDS_PIXELFORMAT),
        ('dwCaps', c_uint32),
        ('dwCaps2', c_uint32),
        ('dwCaps3', c_uint32),
        ('dwCaps4', c_uint32),
        ('dwReversed5', c_uint32)
    ]
    
    _pack_ = True

class DwFlags(IntFlag):
    DDSD_CAPS = 0x1
    DDSD_HEIGHT = 0x2
    DDSD_WIDTH = 0x4
    DDSD_PITCH = 0x8
    DDSD_PIXELFORMAT = 0x1000
    DDSD_MIPMAPCOUNT = 0x20000
    DDSD_LINEARSIZE = 0x80000
    DDSD_DEPTH = 0x800000

class PIXELFORMAT_DwFlags(IntFlag):
    DDPF_ALPHAPIXELS = 0x1
    DDPF_ALPHA = 0x2
    DDPF_FOURCC = 0x4
    DDPF_RGB = 0x40
    DDPF_YUV = 0x200
    DDPF_LUMINANCE = 0x20000


class DwCaps(IntFlag):
    DDSCAPS_COMPLEX = 0x8
    DDSCAPS_MIPMAP = 0x400000
    DDSCAPS_TEXTURE = 0x1000

class DwCaps2(IntFlag):
    DDSCAPS2_CUBEMAP = 0x200
    DDSCAPS2_CUBEMAP_POSITIVEX = 0x400
    DDSCAPS2_CUBEMAP_NEGATIVEX = 0x800
    DDSCAPS2_CUBEMAP_POSITIVEY = 0x1000
    DDSCAPS2_CUBEMAP_NEGATIVEY = 0x2000
    DDSCAPS2_CUBEMAP_POSITIVEZ = 0x4000
    DDSCAPS2_CUBEMAP_NEGATIVEZ = 0x8000
    DDSCAPS2_VOLUME = 0x200000
    
    

def couleur_16bits_vers_rgba(entree : int) -> Tuple[int, int, int, int]:
    
    return (
        int(( entree >> (6 + 5) ) / 0b11111 * 255),
        int(( (entree >> 5) & 0b111111 ) / 0b111111 * 255),
        int(( entree & 0b11111 ) / 0b11111 * 255),
        255
    )


"""
    Notre décodeur se trouve ici. Les dimensions de l'image doivent être
    définies explicitement, puisque autrement on ne sait pas combien de
    blocs sont horizontaux et combien sont verticaux.
"""

def decoder_dxt(version_dxt, hauteur_px : int, largeur_px : int, entree : bytes) -> List[List[Tuple[int, int, int, int]]]:
    
    lecteur = BytesIO(entree)
    
    sortie : List[List[Tuple[int, int, int, int]]] = [[] for colonne in range(hauteur_px + (-hauteur_px % 4))]
    sortie_alpha : List[List[int]] = [[] for colonne in range(hauteur_px + (-hauteur_px % 4))]
    
    
    
    for base_position_y in range(0, hauteur_px, 4):
        
        for base_position_x in range(0, largeur_px, 4):
            
            if version_dxt in (4, 5):
                alpha_1 = lecteur.read(1)[0]
                alpha_2 = lecteur.read(1)[0]
                
                if alpha_1 > alpha_2:
                    alphas = [alpha_1, alpha_2]
                    alphas += [int(alpha_2 * (part/7) + alpha_1 * ((7-part)/7)) for part in range(1, 7)]
                else:
                    alphas = [alpha_1, alpha_2]
                    alphas += [int(alpha_2 * (part/5) + alpha_1 * ((5-part)/5)) for part in range(1, 5)]
                    alphas += [0, 255]
                
                alphas_encodes = int.from_bytes(lecteur.read(6), 'little')
                
                for position_y in range(base_position_y, base_position_y + 4):               
                    for position_x in range(base_position_x, base_position_x + 4):     
                                   
                        sortie_alpha[position_y].append(alphas[alphas_encodes & 0b111])
                        
                        alphas_encodes >>= 3
            
            elif version_dxt in (2, 3):
                alphas_encodes = int.from_bytes(lecteur.read(8))

                for position_y in range(base_position_y, base_position_y + 4):               
                    for position_x in range(base_position_x, base_position_x + 4):     
                                   
                        sortie_alpha[position_y].append(int((alphas_encodes & 0b1111) / 0b1111 * 255))
                        
                        alphas_encodes >>= 4
                
            couleur_1_brute = int.from_bytes(lecteur.read(2), 'little')
            couleur_2_brute = int.from_bytes(lecteur.read(2), 'little')
            
            couleur_1 = couleur_16bits_vers_rgba(couleur_1_brute)
            couleur_2 = couleur_16bits_vers_rgba(couleur_2_brute)
            
            if couleur_1_brute > couleur_2_brute or version_dxt > 1:
                couleur_3 = tuple(
                    int(2/3 * couleur_1[composante] + 1/3 * couleur_2[composante])
                    for composante in range(4)
                )
                couleur_4 = tuple(
                    int(1/3 * couleur_1[composante] + 2/3 * couleur_2[composante])
                    for composante in range(4)
                )
            else:
                couleur_3 = tuple(
                    int(1/2 * couleur_1[composante] + 1/2 * couleur_2[composante])
                    for composante in range(4)
                )
                couleur_4 = (0, 0, 0, 0) # Noir transparent
                
            
            couleurs = [couleur_1, couleur_2, couleur_3, couleur_4]
            
            couleurs_encodees = int.from_bytes(lecteur.read(4), 'little')
            
            decalage_bits = 0
            
            for position_y in range(base_position_y, base_position_y + 4):       
                for position_x in range(base_position_x, base_position_x + 4):              
                
                    if version_dxt == 1:
                        r,g,b,a = couleurs[(couleurs_encodees >> decalage_bits) & 0b11]
                    else:
                        r,g,b = couleurs[(couleurs_encodees >> decalage_bits) & 0b11][:3]
                        a = sortie_alpha[position_y][position_x]
                
                    if version_dxt in (2, 4) and a > 0: # Prémultiplication par l'alpha (dans l'autre sens)
                        r /= (a / 255)
                        b /= (a / 255)
                        g /= (a / 255)
                
                    decalage_bits += 2
                    
                    sortie[position_y].append((int(r), int(g), int(b), int(a)))

    return sortie

def mon_decodeur_dds(entree : bytes):
    
    header = DDS_HEADER()
    lecteur = BytesIO(entree)
    
    lecteur.readinto(header)
    
    print('[INFO] Valeurs de dwFlags => ', DwFlags(header.dwFlags))
    print('[INFO] Valeurs de ddspf.dwFlags => ', PIXELFORMAT_DwFlags(header.ddspf.dwFlags))
    print('[INFO] Valeurs de dwCaps => ', DwCaps(header.dwCaps))
    print('[INFO] Valeurs de dwCaps2 => ', DwCaps(header.dwCaps2))
    
    assert header.dwMagic == b'DDS '
    assert lecteur.tell() - 4 == header.dwSize == 124
    
    image = Image.new('RGBA', (header.dwWidth, header.dwHeight))
    
    if header.ddspf.dwFourCC[:3] == b'DXT':
        version_dxt = int(header.ddspf.dwFourCC[3:])
        
        assert 1 <= version_dxt <= 5
        
        couleurs_rgba : List[List[Tuple[int, int, int, int]]] = decoder_dxt(version_dxt, header.dwHeight, header.dwWidth, lecteur.read())
    
    else:
        raise NotImplementedError('Format de compression non supporté : ' + str(header.ddspf.dwFourCC))

    image_pixels = image.load()

    for y in range(header.dwHeight):
        for x in range(header.dwWidth):
            image_pixels[x, y] = couleurs_rgba[y][x]
    
    image.show() # L'image est affichée directement dans une nouvelle fenêtre !

# Prendre un chemin de la part de l'utilisateur en entrée.

args = ArgumentParser(description = "Affiche une texture DDS de petite taille " +
    "à l'écran, en utilisant la bibliothèque PIL.")

args.add_argument('fichier_entree', help = "Chemin de l'image à décoder sur le disque.")

args = args.parse_args()

with open(args.fichier_entree, 'rb') as objet_fichier:
    mon_decodeur_dds(objet_fichier.read())