PlayStation 1 .TIM/.PXL/.CLT to PNG image converter (with support for multiple palettes)

tim_to_png.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""PlayStation 1 .TIM/.PXL/.CLT to PNG image converter

A command-line tool to decode an image file in the .TIM format commonly used on
the PS1, with proper support for images with multiple color palettes (a separate
PNG file will be generated for each palette available) as well as the somewhat
uncommon split .PXL (image) + .CLT (palette) format. Requires PIL/Pillow and
NumPy to be installed.
"""

__version__ = "0.1.0"
__author__  = "spicyjpeg"

from argparse        import ArgumentParser, Namespace
from collections.abc import Generator
from dataclasses     import dataclass, field
from enum            import IntEnum, IntFlag
from pathlib         import Path
from struct          import Struct
from typing          import Self

import numpy
from numpy import ndarray
from PIL   import Image

## Utilities

TRANSPARENT_COLOR: int = 0x0000

def toRGBA(inputData: ndarray) -> ndarray:
	source: ndarray = inputData.astype("<H")
	r:      ndarray = (source >>  0) & 31
	g:      ndarray = (source >>  5) & 31
	b:      ndarray = (source >> 10) & 31

	alpha: ndarray = numpy.full_like(source, 255)

	numpy.copyto(alpha, 128, where = ((inputData >> 15) == 1))
	numpy.copyto(alpha,   0, where = (inputData == TRANSPARENT_COLOR))

	return numpy.stack((
		((r * 255) + 15) // 31,
		((g * 255) + 15) // 31,
		((b * 255) + 15) // 31,
		alpha
	), 2, dtype = "B")

def slice2DArray(
	data:       ndarray,
	tileWidth:  int,
	tileHeight: int
) -> Generator[ndarray, None, None]:
	if (data.shape[1] % tileWidth) or (data.shape[0] % tileHeight):
		raise ValueError(
			f"input array dimensions are not an exact multiple of "
			f"{tileWidth}x{tileHeight}"
		)

	for y in range(0, data.shape[0], tileHeight):
		for x in range(0, data.shape[1], tileWidth):
			yield data[y:y + tileHeight, x:x + tileWidth]

def fixImagePalette(imageObj: Image.Image):
	if imageObj.palette.mode != "RGBA":
		return

	clut: ndarray = numpy.array(imageObj.getpalette("RGBA"), "B")
	clut          = clut.reshape(( clut.shape[0] // 4, 4 ))

	# Pillow's PNG decoder does not handle indexed color images with alpha
	# correctly, so a workaround is needed here to manually split off the alpha
	# channel into the image's "tRNs" chunk.
	# TODO: this seems to be no longer needed...?
	imageObj.putpalette(clut[:, 0:2].tobytes(), "RGB")
	imageObj.info["transparency"] = clut[:, 3].tobytes().rstrip(b"\xff")

## .TIM file parser

class TIMColorDepth(IntEnum):
	COLOR_4BPP  = 0
	COLOR_8BPP  = 1
	COLOR_16BPP = 2

class TIMHeaderFlag(IntFlag):
	COLOR_BITMASK = 3 << 0
	HAS_PALETTE   = 1 << 3

TIM_HEADER_STRUCT:  Struct = Struct("< 2I")
TIM_SECTION_STRUCT: Struct = Struct("< I 4H")
TIM_HEADER_VERSION: int    = 0x10
CLT_HEADER_VERSION: int    = 0x11
PXL_HEADER_VERSION: int    = 0x12

@dataclass
class TIMSection:
	x:    int
	y:    int
	data: ndarray = field(repr = False)

	@staticmethod
	def parse(
		data:      bytes | bytearray,
		offset:    int = 0,
		pixelSize: int = 2
	) -> tuple[int, Self]:
		(
			endOffset,
			x,
			y,
			width,
			height
		) = TIM_SECTION_STRUCT.unpack_from(data, offset)

		data    = data[offset + TIM_SECTION_STRUCT.size:offset + endOffset]
		offset += endOffset

		if len(data) != (width * height * 2):
			raise RuntimeError("section length does not match image size")

		image: ndarray = numpy.frombuffer(
			data,
			"<H" if (pixelSize == 2) else "B"
		).reshape((
			height,
			(width * 2) // pixelSize
		))

		return offset, TIMSection(x, y, image)

@dataclass
class TIMImage:
	colorDepth: TIMColorDepth
	image:      TIMSection | None = None
	clut:       TIMSection | None = None

	@staticmethod
	def parse(data: bytes | bytearray, offset: int = 0) -> Self:
		version, flags = TIM_HEADER_STRUCT.unpack_from(data, offset)
		offset        += TIM_HEADER_STRUCT.size

		if version == TIM_HEADER_VERSION:
			hasImage: bool = True
			hasCLUT:  bool = bool(flags & TIMHeaderFlag.HAS_PALETTE)
		else:
			hasImage: bool = (version == PXL_HEADER_VERSION)
			hasCLUT:  bool = (version == CLT_HEADER_VERSION)

			if not (hasImage or hasCLUT):
				raise ValueError(f"invalid .TIM file version: {version:#x}")

		colorDepth: TIMColorDepth = \
			TIMColorDepth(flags & TIMHeaderFlag.COLOR_BITMASK)
		pixelSize:  int           = \
			2 if (colorDepth == TIMColorDepth.COLOR_16BPP) else 1

		tim: TIMImage = TIMImage(colorDepth)

		if hasCLUT:
			offset, tim.clut  = TIMSection.parse(data, offset, 2)
		if hasImage:
			offset, tim.image = TIMSection.parse(data, offset, pixelSize)

		return tim

	def _getImageData(self) -> ndarray:
		match self.colorDepth:
			case TIMColorDepth.COLOR_4BPP:
				# Unpack each byte into two adjacent pixels for 4bpp images.
				data: ndarray = numpy.zeros((
					self.image.data.shape[0],
					self.image.data.shape[1] * 2
				), "B")
				data[:, 0::2] = self.image.data & 15
				data[:, 1::2] = self.image.data >> 4

				return data

			case TIMColorDepth.COLOR_8BPP:
				return self.image.data

			case TIMColorDepth.COLOR_16BPP:
				return toRGBA(self.image.data)

	def toImage(self) -> Image.Image:
		data: ndarray = self._getImageData()

		if self.clut is None:
			image: Image.Image = Image.fromarray(data, "RGBA")
		else:
			image: Image.Image = Image.fromarray(data, "P")

			image.putpalette(toRGBA(self.clut.data), "RGBA")

		return image

	def toGrayscaleImage(self) -> Image.Image:
		if self.colorDepth == TIMColorDepth.COLOR_16BPP:
			raise RuntimeError(
				"grayscale conversion is only supported for 4bpp and 8bpp "
				"images"
			)

		data: ndarray = self._getImageData()

		# Scale the brightness from 0-15 to 0-255 range for 4bpp images.
		if self.colorDepth == TIMColorDepth.COLOR_4BPP:
			data *= 17

		return Image.fromarray(data, "L")

## Main

def createParser() -> ArgumentParser:
	parser = ArgumentParser(
		description = \
			"Decodes a PlayStation 1 .TIM or .PXL image file, optionally "
			"alongside one or more .CLT (color palette) files, and outputs a "
			"PNG file for each image and palette combination.",
		add_help    = False
	)

	group = parser.add_argument_group("Tool options")
	group.add_argument(
		"-h", "--help",
		action = "help",
		help   = "Show this help message and exit"
	)

	group = parser.add_argument_group("Conversion options")
	group.add_argument(
		"-r", "--rgba",
		action = "store_true",
		help   = \
			"Output PNG files in RGBA format with no palette instead of "
			"indexed color"
	)
	group.add_argument(
		"-g", "--grayscale",
		action = "store_true",
		help   = \
			"Generate an additional copy of the image with a dummy grayscale "
			"palette applied"
	)
	group.add_argument(
		"-i", "--ignore-tim-palette",
		action = "store_true",
		help   = \
			"Ignore any palettes embedded within the .TIM file (only use "
			"palettes from the specified .CLT files and/or the grayscale "
			"palette if -g is used)"
	)
	group.add_argument(
		"-u", "--skip-unlikely",
		action = "store_true",
		help   = \
			""
			"palettes from the specified .CLT files and/or the grayscale "
			"palette if -g is used)"
	)

	group = parser.add_argument_group("File paths")
	group.add_argument(
		"-o", "--output",
		type    = Path,
		help    =
			"Save generated PNG files to given path (same directory as input "
			"image by default)",
		metavar = "dir"
	)
	group.add_argument(
		"image",
		type = Path,
		help = "Path to .TIM or .PXL file to convert"
	)
	group.add_argument(
		"palette",
		type  = Path,
		nargs = "*",
		help  = \
			"Paths to additional palettes (.CLT files) for the image; at least "
			"one must be specified if the input image is an 4bpp or 8bpp .PXL "
			"file and -g is not used"
	)

	return parser

def main():
	parser: ArgumentParser = createParser()
	args:   Namespace      = parser.parse_args()

	with open(args.image, "rb") as file:
		tim: TIMImage = TIMImage.parse(file.read())

	# Gather all palettes into sets (one per .TIM/.CLT file, plus one for the
	# default grayscale palettes).
	paletteSets: dict[str, list[ndarray]] = {}
	numColors:   int                      = \
		256 if (tim.colorDepth == TIMColorDepth.COLOR_8BPP) else 16

	if not args.ignore_tim_palette and (tim.clut is not None):
		slices: Generator = slice2DArray(tim.clut.data, numColors, 1)
		paletteSets[""]   = list(slices)

	for path in args.palette:
		with open(path, "rb") as file:
			palette: TIMImage = TIMImage.parse(file.read())

		# If the image is 4bpp but the palette data's width is 256 pixels or a
		# multiple thereof, assume the palette is for an 8bpp image (rather than
		# 16 4bpp palettes laid out horizontally) and skip it.
		if palette.clut.data.shape[1] < numColors:
			continue
		if (1
	  		and args.skip_unlikely
	  		and (tim.colorDepth == TIMColorDepth.COLOR_4BPP)
	  		and not (palette.clut.data.shape[1] % 256)
		):
			continue

		slices: Generator      = slice2DArray(palette.clut.data, numColors, 1)
		paletteSets[path.stem] = list(slices)

	is16bpp:     bool = (tim.colorDepth == TIMColorDepth.COLOR_16BPP)
	hasPalettes: bool = bool(paletteSets or args.grayscale)

	if not is16bpp and not hasPalettes:
		parser.error(
			"input image has no embedded palettes; at least one palette file "
			"must be provided, or the -g option passed to use a grayscale "
			"palette"
		)
	if is16bpp and hasPalettes:
		parser.error("input image is 16bpp, no palettes should be provided")

	# Go through each palette set and generate the respective PNG files.
	outputPath: Path = args.output or args.image.parent

	if args.grayscale:
		image: Image.Image = tim.toGrayscaleImage()

		image.save(outputPath / f"{args.image.stem}-gray.png")

	for setName, palettes in paletteSets.items():
		for index, clut in enumerate(palettes):
			name: str = args.image.stem

			if (len(paletteSets) > 1) and setName:
				name += f"-{setName}"
			if len(palettes) > 1:
				name += f"-{index}"

			tim.clut           = TIMSection(0, 0, clut)
			image: Image.Image = tim.toImage()

			if args.rgba:
				image = image.convert("RGBA")
			#else:
				#fixImagePalette(image)

			image.save(outputPath / f"{name}.png")

if __name__ == "__main__":
	main()