mirror of
https://github.com/KrisKennaway/ii-pix.git
synced 2024-11-19 23:32:18 +00:00
453 lines
19 KiB
Python
453 lines
19 KiB
Python
"""Image converter to Apple II Double Hi-Res format."""
|
|
|
|
import argparse
|
|
import os.path
|
|
from typing import Tuple, List
|
|
|
|
from PIL import Image
|
|
import colour
|
|
import numpy as np
|
|
from sklearn import cluster
|
|
|
|
from os import environ
|
|
|
|
environ['PYGAME_HIDE_SUPPORT_PROMPT'] = '1'
|
|
import pygame
|
|
|
|
import dither as dither_pyx
|
|
import dither_pattern
|
|
import image as image_py
|
|
import palette as palette_py
|
|
import screen as screen_py
|
|
|
|
|
|
# TODO:
|
|
# - support LR/DLR
|
|
# - support HGR
|
|
|
|
|
|
class ClusterPalette:
|
|
def __init__(
|
|
self, image: Image, rgb12_iigs_to_cam16ucs, rgb24_to_cam16ucs,
|
|
reserved_colours=0):
|
|
self._image_rgb = image
|
|
self._colours_cam = self._image_colours_cam(image)
|
|
|
|
self._errors = [1e9] * 16
|
|
|
|
# We fit a 16-colour palette against the entire image which is used
|
|
# as starting values for fitting the 16 SHR palettes. This helps to
|
|
# provide better global consistency of colours across the palettes,
|
|
# e.g. for large blocks of colour. Otherwise these can take a while
|
|
# to converge.
|
|
self._global_palette = np.empty((16, 3), dtype=np.uint8)
|
|
|
|
# How many image colours to fix identically across all 16 SHR
|
|
# palettes. These are taken to be the most prevalent colours from
|
|
# _global_palette.
|
|
self._reserved_colours = reserved_colours
|
|
|
|
# 16 SHR palettes each of 16 colours, in CAM16UCS format
|
|
self._palettes_cam = np.empty((16, 16, 3), dtype=np.float32)
|
|
|
|
# 16 SHR palettes each of 16 colours, in //gs 4-bit RGB format
|
|
self._palettes_rgb = np.empty((16, 16, 3), dtype=np.uint8)
|
|
|
|
# Conversion matrix from 12-bit //gs RGB colour space to CAM16UCS
|
|
# colour space
|
|
self._rgb12_iigs_to_cam16ucs = rgb12_iigs_to_cam16ucs
|
|
|
|
self._rgb24_to_cam16ucs = rgb24_to_cam16ucs
|
|
|
|
# List of line ranges used to train the 16 SHR palettes
|
|
# [(lower_0, upper_0), ...]
|
|
self._palette_splits = self._equal_palette_splits()
|
|
|
|
# Whether the previous iteration of proposed palettes was accepted
|
|
self._palettes_accepted = False
|
|
|
|
# Which palette index's line ranges did we mutate in previous iteration
|
|
self._palette_mutate_idx = 0
|
|
|
|
# Delta applied to palette split in previous iteration
|
|
self._palette_mutate_delta = (0, 0)
|
|
|
|
def _image_colours_cam(self, image: Image):
|
|
colours_rgb = np.asarray(image).reshape((-1, 3))
|
|
with colour.utilities.suppress_warnings(colour_usage_warnings=True):
|
|
colours_cam = colour.convert(colours_rgb, "RGB",
|
|
"CAM16UCS").astype(np.float32)
|
|
return colours_cam
|
|
|
|
def _equal_palette_splits(self, palette_height=35):
|
|
# The 16 palettes are striped across consecutive (overlapping) line
|
|
# ranges. Since nearby lines tend to have similar colours, this has
|
|
# the effect of smoothing out the colour transitions across palettes.
|
|
|
|
# If we want to overlap 16 palettes in 200 lines, where each palette
|
|
# has height H and overlaps the previous one by L lines, then the
|
|
# boundaries are at lines:
|
|
# (0, H), (H-L, 2H-L), (2H-2L, 3H-2L), ..., (15H-15L, 16H - 15L)
|
|
# i.e. 16H - 15L = 200, so for a given palette height H we need to
|
|
# overlap by:
|
|
# L = (16H - 200)/15
|
|
|
|
palette_overlap = (16 * palette_height - 200) / 15
|
|
|
|
palette_ranges = []
|
|
for palette_idx in range(16):
|
|
palette_lower = palette_idx * (palette_height - palette_overlap)
|
|
palette_upper = palette_lower + palette_height
|
|
palette_ranges.append((int(np.round(palette_lower)),
|
|
int(np.round(palette_upper))))
|
|
return palette_ranges
|
|
|
|
def _dither_image(self, palettes_cam, penalty):
|
|
# Suppress divide by zero warning,
|
|
# https://github.com/colour-science/colour/issues/900
|
|
with colour.utilities.suppress_warnings(python_warnings=True):
|
|
palettes_linear_rgb = colour.convert(
|
|
palettes_cam, "CAM16UCS", "RGB").astype(np.float32)
|
|
|
|
output_4bit, line_to_palette, total_image_error = \
|
|
dither_pyx.dither_shr(
|
|
self._image_rgb, palettes_cam, palettes_linear_rgb,
|
|
self._rgb24_to_cam16ucs, float(penalty))
|
|
|
|
return (output_4bit, line_to_palette, palettes_linear_rgb,
|
|
total_image_error)
|
|
|
|
def iterate(self, penalty: float, max_iterations: int):
|
|
iterations_since_improvement = 0
|
|
total_image_error = 1e9
|
|
last_good_splits = self._palette_splits
|
|
|
|
while iterations_since_improvement < max_iterations:
|
|
# print("Iterations %d" % iterations_since_improvement)
|
|
new_palettes_cam, new_palettes_rgb12_iigs, new_palette_errors = (
|
|
self._propose_palettes())
|
|
|
|
# Recompute image with proposed palettes and check whether it has
|
|
# lower total image error than our previous best.
|
|
(output_4bit, line_to_palette, palettes_linear_rgb,
|
|
new_total_image_error) = self._dither_image(
|
|
new_palettes_cam, penalty)
|
|
|
|
self._reassign_unused_palettes(line_to_palette,
|
|
last_good_splits)
|
|
|
|
if new_total_image_error >= total_image_error:
|
|
iterations_since_improvement += 1
|
|
continue
|
|
|
|
# We found a globally better set of palettes
|
|
iterations_since_improvement = 0
|
|
last_good_splits = self._palette_splits
|
|
total_image_error = new_total_image_error
|
|
|
|
self._palettes_cam = new_palettes_cam
|
|
self._palettes_rgb = new_palettes_rgb12_iigs
|
|
self._errors = new_palette_errors
|
|
self._palettes_accepted = True
|
|
|
|
yield (new_total_image_error, output_4bit, line_to_palette,
|
|
new_palettes_rgb12_iigs, palettes_linear_rgb)
|
|
|
|
def _propose_palettes(self) -> Tuple[np.ndarray, np.ndarray, List[float]]:
|
|
"""Attempt to find new palettes that locally improve image quality.
|
|
|
|
Re-fit a set of 16 palettes from (overlapping) line ranges of the
|
|
source image, using k-means clustering in CAM16-UCS colour space.
|
|
|
|
We maintain the total image error for the pixels on which the 16
|
|
palettes are clustered. A new palette that increases this local
|
|
image error is rejected.
|
|
|
|
New palettes that reduce local error cannot be applied immediately
|
|
though, because they may cause an increase in *global* image error
|
|
when dithering. i.e. they would reduce the overall image quality.
|
|
|
|
The current (locally) best palettes are returned and can be applied
|
|
using accept_palettes().
|
|
"""
|
|
new_errors = list(self._errors)
|
|
new_palettes_cam = np.empty_like(self._palettes_cam)
|
|
new_palettes_rgb12_iigs = np.empty_like(self._palettes_rgb)
|
|
|
|
# Compute a new 16-colour global palette for the entire image,
|
|
# used as the starting center positions for k-means clustering of the
|
|
# individual palettes
|
|
self._fit_global_palette()
|
|
|
|
self._mutate_palette_splits()
|
|
for palette_idx in range(16):
|
|
palette_lower, palette_upper = self._palette_splits[palette_idx]
|
|
palette_pixels = self._colours_cam[
|
|
palette_lower * 320:palette_upper * 320, :]
|
|
|
|
palettes_rgb12_iigs, palette_error = \
|
|
dither_pyx.k_means_with_fixed_centroids(
|
|
n_clusters=16, n_fixed=self._reserved_colours,
|
|
samples=palette_pixels,
|
|
initial_centroids=self._global_palette,
|
|
max_iterations=1000, tolerance=0.05,
|
|
rgb12_iigs_to_cam16ucs=self._rgb12_iigs_to_cam16ucs
|
|
)
|
|
|
|
if (palette_error >= self._errors[palette_idx] and not
|
|
self._reserved_colours):
|
|
# Not a local improvement to the existing palette, so ignore it.
|
|
# We can't take this shortcut when we're reserving colours
|
|
# because it would break the invariant that all palettes must
|
|
# share colours.
|
|
continue
|
|
for i in range(16):
|
|
new_palettes_cam[palette_idx, i, :] = (
|
|
np.array(dither_pyx.convert_rgb12_iigs_to_cam(
|
|
self._rgb12_iigs_to_cam16ucs, palettes_rgb12_iigs[
|
|
i]), dtype=np.float32))
|
|
|
|
new_palettes_rgb12_iigs[palette_idx, :, :] = palettes_rgb12_iigs
|
|
new_errors[palette_idx] = palette_error
|
|
|
|
self._palettes_accepted = False
|
|
return new_palettes_cam, new_palettes_rgb12_iigs, new_errors
|
|
|
|
def _fit_global_palette(self):
|
|
"""Compute a 16-colour palette for the entire image to use as
|
|
starting point for the sub-palettes. This should help when the image
|
|
has large blocks of colour since the sub-palettes will tend to pick the
|
|
same colours."""
|
|
|
|
clusters = cluster.MiniBatchKMeans(n_clusters=16, max_iter=10000)
|
|
clusters.fit_predict(self._colours_cam)
|
|
|
|
# Dict of {palette idx : frequency count}
|
|
palette_freq = {idx: 0 for idx in range(16)}
|
|
for idx, freq in zip(*np.unique(clusters.labels_, return_counts=True)):
|
|
palette_freq[idx] = freq
|
|
frequency_order = [
|
|
k for k, v in sorted(
|
|
list(palette_freq.items()), key=lambda kv: kv[1], reverse=True)]
|
|
|
|
self._global_palette = (
|
|
dither_pyx.convert_cam16ucs_to_rgb12_iigs(
|
|
clusters.cluster_centers_[frequency_order].astype(
|
|
np.float32)))
|
|
|
|
def _mutate_palette_splits(self):
|
|
if self._palettes_accepted:
|
|
# Last time was good, keep going
|
|
self._apply_palette_delta(self._palette_mutate_idx,
|
|
self._palette_mutate_delta[0],
|
|
self._palette_mutate_delta[1])
|
|
else:
|
|
# undo last mutation
|
|
self._apply_palette_delta(self._palette_mutate_idx,
|
|
-self._palette_mutate_delta[0],
|
|
-self._palette_mutate_delta[1])
|
|
|
|
# Pick a palette endpoint to move up or down
|
|
palette_to_mutate = np.random.randint(0, 16)
|
|
while True:
|
|
if palette_to_mutate > 0:
|
|
palette_lower_delta = np.random.randint(-20, 21)
|
|
else:
|
|
palette_lower_delta = 0
|
|
if palette_to_mutate < 15:
|
|
palette_upper_delta = np.random.randint(-20, 21)
|
|
else:
|
|
palette_upper_delta = 0
|
|
if palette_lower_delta != 0 or palette_upper_delta != 0:
|
|
break
|
|
|
|
self._apply_palette_delta(palette_to_mutate, palette_lower_delta,
|
|
palette_upper_delta)
|
|
|
|
def _apply_palette_delta(
|
|
self, palette_to_mutate, palette_lower_delta, palette_upper_delta):
|
|
old_lower, old_upper = self._palette_splits[palette_to_mutate]
|
|
new_lower = old_lower + palette_lower_delta
|
|
new_upper = old_upper + palette_upper_delta
|
|
|
|
new_lower = np.clip(new_lower, 0, np.clip(new_upper, 1, 200) - 1)
|
|
new_upper = np.clip(new_upper, new_lower + 1, 200)
|
|
assert new_lower >= 0, new_upper - 1
|
|
|
|
self._palette_splits[palette_to_mutate] = (new_lower, new_upper)
|
|
self._palette_mutate_idx = palette_to_mutate
|
|
self._palette_mutate_delta = (palette_lower_delta, palette_upper_delta)
|
|
|
|
def _reassign_unused_palettes(self, new_line_to_palette, last_good_splits):
|
|
palettes_used = [False] * 16
|
|
for palette in new_line_to_palette:
|
|
palettes_used[palette] = True
|
|
for palette_idx, palette_used in enumerate(palettes_used):
|
|
if palette_used:
|
|
continue
|
|
print("Reassigning palette %d" % palette_idx)
|
|
max_width = 0
|
|
split_palette_idx = -1
|
|
idx = 0
|
|
for lower, upper in last_good_splits:
|
|
width = upper - lower
|
|
if width > max_width:
|
|
split_palette_idx = idx
|
|
idx += 1
|
|
|
|
lower, upper = last_good_splits[split_palette_idx]
|
|
if upper - lower > 20:
|
|
mid = (lower + upper) // 2
|
|
self._palette_splits[split_palette_idx] = (
|
|
lower, mid - 1)
|
|
self._palette_splits[palette_idx] = (mid, upper)
|
|
else:
|
|
lower = np.random.randint(0, 199)
|
|
upper = np.random.randint(lower + 1, 200)
|
|
self._palette_splits[palette_idx] = (lower, upper)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("input", type=str, help="Input image file to process.")
|
|
parser.add_argument("output", type=str, help="Output file for converted "
|
|
"Apple II image.")
|
|
parser.add_argument(
|
|
"--lookahead", type=int, default=8,
|
|
help=("How many pixels to look ahead to compensate for NTSC colour "
|
|
"artifacts (default: 8)"))
|
|
parser.add_argument(
|
|
'--dither', type=str, choices=list(dither_pattern.PATTERNS.keys()),
|
|
default=dither_pattern.DEFAULT_PATTERN,
|
|
help="Error distribution pattern to apply when dithering (default: "
|
|
+ dither_pattern.DEFAULT_PATTERN + ")")
|
|
parser.add_argument(
|
|
'--show-input', action=argparse.BooleanOptionalAction, default=False,
|
|
help="Whether to show the input image before conversion.")
|
|
parser.add_argument(
|
|
'--show-output', action=argparse.BooleanOptionalAction, default=True,
|
|
help="Whether to show the output image after conversion.")
|
|
parser.add_argument(
|
|
'--palette', type=str, choices=list(set(palette_py.PALETTES.keys())),
|
|
default=palette_py.DEFAULT_PALETTE,
|
|
help='RGB colour palette to dither to. "ntsc" blends colours over 8 '
|
|
'pixels and gives better image quality on targets that '
|
|
'use/emulate NTSC, but can be substantially slower. Other '
|
|
'palettes determine colours based on 4 pixel sequences '
|
|
'(default: ' + palette_py.DEFAULT_PALETTE + ")")
|
|
parser.add_argument(
|
|
'--show-palette', type=str, choices=list(palette_py.PALETTES.keys()),
|
|
help="RGB colour palette to use when --show_output (default: "
|
|
"value of --palette)")
|
|
parser.add_argument(
|
|
'--verbose', action=argparse.BooleanOptionalAction,
|
|
default=False, help="Show progress during conversion")
|
|
parser.add_argument(
|
|
'--gamma_correct', type=float, default=2.4,
|
|
help='Gamma-correct image by this value (default: 2.4)'
|
|
)
|
|
args = parser.parse_args()
|
|
if args.lookahead < 1:
|
|
parser.error('--lookahead must be at least 1')
|
|
|
|
# palette = palette_py.PALETTES[args.palette]()
|
|
screen = screen_py.SHR320Screen()
|
|
|
|
# Conversion matrix from RGB to CAM16UCS colour values. Indexed by
|
|
# 24-bit RGB value
|
|
rgb24_to_cam16ucs = np.load("data/rgb24_to_cam16ucs.npy")
|
|
rgb12_iigs_to_cam16ucs = np.load("data/rgb12_iigs_to_cam16ucs.npy")
|
|
|
|
# Open and resize source image
|
|
image = image_py.open(args.input)
|
|
if args.show_input:
|
|
image_py.resize(image, screen.X_RES, screen.Y_RES,
|
|
srgb_output=False).show()
|
|
rgb = np.array(
|
|
image_py.resize(image, screen.X_RES, screen.Y_RES,
|
|
gamma=args.gamma_correct)).astype(np.float32) / 255
|
|
|
|
# TODO: flags
|
|
penalty = 1 # 1e18 # TODO: is this needed any more?
|
|
iterations = 200
|
|
|
|
pygame.init()
|
|
# TODO: for some reason I need to execute this twice - the first time
|
|
# the window is created and immediately destroyed
|
|
_ = pygame.display.set_mode((640, 400))
|
|
canvas = pygame.display.set_mode((640, 400))
|
|
canvas.fill((0, 0, 0))
|
|
pygame.display.flip()
|
|
|
|
total_image_error = None
|
|
# TODO: reserved_colours should be a flag
|
|
cluster_palette = ClusterPalette(
|
|
rgb, reserved_colours=1,
|
|
rgb12_iigs_to_cam16ucs=rgb12_iigs_to_cam16ucs,
|
|
rgb24_to_cam16ucs=rgb24_to_cam16ucs)
|
|
|
|
for (new_total_image_error, output_4bit, line_to_palette,
|
|
palettes_rgb12_iigs, palettes_linear_rgb) in cluster_palette.iterate(
|
|
penalty, iterations):
|
|
|
|
if total_image_error is not None:
|
|
print("Improved quality +%f%% (%f)" % (
|
|
(1 - new_total_image_error / total_image_error) * 100,
|
|
new_total_image_error))
|
|
total_image_error = new_total_image_error
|
|
for i in range(16):
|
|
screen.set_palette(i, palettes_rgb12_iigs[i, :, :])
|
|
|
|
# Recompute current screen RGB image
|
|
screen.set_pixels(output_4bit)
|
|
output_rgb = np.empty((200, 320, 3), dtype=np.uint8)
|
|
for i in range(200):
|
|
screen.line_palette[i] = line_to_palette[i]
|
|
output_rgb[i, :, :] = (
|
|
palettes_linear_rgb[line_to_palette[i]][
|
|
output_4bit[i, :]] * 255
|
|
).astype(np.uint8)
|
|
|
|
output_srgb = (image_py.linear_to_srgb(output_rgb)).astype(np.uint8)
|
|
|
|
# dither = dither_pattern.PATTERNS[args.dither]()
|
|
# bitmap = dither_pyx.dither_image(
|
|
# screen, rgb, dither, args.lookahead, args.verbose, rgb24_to_cam16ucs)
|
|
|
|
# Show output image by rendering in target palette
|
|
# output_palette_name = args.show_palette or args.palette
|
|
# output_palette = palette_py.PALETTES[output_palette_name]()
|
|
# output_screen = screen_py.DHGRScreen(output_palette)
|
|
# if output_palette_name == "ntsc":
|
|
# output_srgb = output_screen.bitmap_to_image_ntsc(bitmap)
|
|
# else:
|
|
# output_srgb = image_py.linear_to_srgb(
|
|
# output_screen.bitmap_to_image_rgb(bitmap)).astype(np.uint8)
|
|
out_image = image_py.resize(
|
|
Image.fromarray(output_srgb), screen.X_RES * 2, screen.Y_RES * 2,
|
|
srgb_output=True)
|
|
|
|
if args.show_output:
|
|
surface = pygame.surfarray.make_surface(
|
|
np.asarray(out_image).transpose((1, 0, 2))) # flip y/x axes
|
|
canvas.blit(surface, (0, 0))
|
|
pygame.display.flip()
|
|
# print((palettes_rgb * 255).astype(np.uint8))
|
|
unique_colours = np.unique(
|
|
palettes_rgb12_iigs.reshape(-1, 3), axis=0).shape[0]
|
|
print("%d unique colours" % unique_colours)
|
|
|
|
# Save Double hi-res image
|
|
outfile = os.path.join(os.path.splitext(args.output)[0] + "-preview.png")
|
|
out_image.save(outfile, "PNG")
|
|
screen.pack()
|
|
# with open(args.output, "wb") as f:
|
|
# f.write(bytes(screen.aux))
|
|
# f.write(bytes(screen.main))
|
|
with open(args.output, "wb") as f:
|
|
f.write(bytes(screen.memory))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|