mirror of
https://github.com/KrisKennaway/ii-vision.git
synced 2024-09-29 18:59:05 +00:00
eebbccf711
Clean up naming in edit_distance In video encoder, when we emit additional offsets as part of an opcode, reinsert back into the priority heapq if the new edit distance is nonzero, in case we get the chance to fix it up later in the frame. Also make sure to zero out the diff_weights and content_deltas so we don't consider the offset again as a side-effect of some other opcode. Instead of prioritizing side-effect offsets by their previous update priority, prioritize by those with the lowest (error - edit) delta i.e. not introducing too much error relative to their edit distance.
308 lines
10 KiB
Python
308 lines
10 KiB
Python
"""Encode a sequence of images as an optimized stream of screen changes."""
|
|
|
|
import heapq
|
|
import os
|
|
import queue
|
|
import random
|
|
import subprocess
|
|
import threading
|
|
from typing import List, Iterator, Tuple
|
|
|
|
# import hitherdither
|
|
import numpy as np
|
|
import skvideo.io
|
|
from PIL import Image
|
|
|
|
import edit_distance
|
|
import opcodes
|
|
import screen
|
|
|
|
|
|
class Video:
|
|
"""Apple II screen memory map encoding a bitmapped frame."""
|
|
|
|
CLOCK_SPEED = 1024 * 1024 # type: int
|
|
|
|
def __init__(
|
|
self,
|
|
filename: str,
|
|
):
|
|
self.filename = filename # type: str
|
|
|
|
self._reader = skvideo.io.FFmpegReader(filename)
|
|
|
|
# Compute frame rate from input video
|
|
# TODO: possible to compute time offset for each frame instead?
|
|
data = skvideo.io.ffprobe(self.filename)['video']
|
|
rate_data = data['@r_frame_rate'].split("/") # e.g. 12000/1001
|
|
self._input_frame_rate = float(
|
|
rate_data[0]) / float(rate_data[1]) # type: float
|
|
|
|
self.cycles_per_frame = (
|
|
self.CLOCK_SPEED / self._input_frame_rate) # type: float
|
|
self.frame_number = 0 # type: int
|
|
|
|
# Initialize empty screen
|
|
self.memory_map = screen.MemoryMap(
|
|
screen_page=1) # type: screen.MemoryMap
|
|
|
|
# Accumulates pending edit weights across frames
|
|
self.update_priority = np.zeros((32, 256), dtype=np.int64)
|
|
|
|
def tick(self, cycles: int) -> bool:
|
|
if cycles > (self.cycles_per_frame * self.frame_number):
|
|
self.frame_number += 1
|
|
return True
|
|
return False
|
|
|
|
def _frame_grabber(self) -> Iterator[Image.Image]:
|
|
for frame_array in self._reader.nextFrame():
|
|
yield Image.fromarray(frame_array)
|
|
|
|
@staticmethod
|
|
def _rgb(r, g, b):
|
|
return (r << 16) + (g << 8) + b
|
|
|
|
# def dither_framesframes(self) -> Iterator[screen.MemoryMap]:
|
|
# palette = hitherdither.palette.Palette(
|
|
# [
|
|
# self._rgb(0,0,0), # black */
|
|
# self._rgb(148,12,125), # red - hgr 0*/
|
|
# self._rgb(32,54,212), # dk blue - hgr 0 */
|
|
# self._rgb(188,55,255), # purple - default HGR overlay color */
|
|
# self._rgb(51,111,0), # dk green - hgr 0 */
|
|
# self._rgb(126,126,126), # gray - hgr 0 */
|
|
# self._rgb(7,168,225), # med blue - alternate HGR overlay
|
|
# # color */
|
|
# self._rgb(158,172,255), # lt blue - hgr 0 */
|
|
# self._rgb(99,77,0), # brown - hgr 0 */
|
|
# self._rgb(249,86,29), # orange */
|
|
# self._rgb(126,126,126), # grey - hgr 0 */
|
|
# self._rgb(255,129,236), # pink - hgr 0 */
|
|
# self._rgb(67,200,0), # lt green */
|
|
# self._rgb(221,206,23), # yellow - hgr 0 */
|
|
# self._rgb(93,248,133), # aqua - hgr 0 */
|
|
# self._rgb(255,255,255) # white
|
|
# ]
|
|
# )
|
|
# for _idx, _frame in enumerate(self._frame_grabber()):
|
|
# if _idx % 60 == 0:
|
|
# img_dithered = hitherdither.ordered.yliluoma.yliluomas_1_ordered_dithering(
|
|
# _frame.resize((280,192), resample=Image.NEAREST),
|
|
# palette, order=8)
|
|
#
|
|
# yield img_dithered
|
|
|
|
def frames(self) -> Iterator[screen.MemoryMap]:
|
|
"""Encode frame to HGR using bmp2dhr.
|
|
|
|
We do the encoding in a background thread to parallelize.
|
|
"""
|
|
|
|
frame_dir = self.filename.split(".")[0]
|
|
try:
|
|
os.mkdir(frame_dir)
|
|
except FileExistsError:
|
|
pass
|
|
|
|
q = queue.Queue(maxsize=10)
|
|
|
|
def worker():
|
|
"""Invoke bmp2dhr to encode input image frames and push to queue."""
|
|
for _idx, _frame in enumerate(self._frame_grabber()):
|
|
outfile = "%s/%08dC.BIN" % (frame_dir, _idx)
|
|
bmpfile = "%s/%08d.bmp" % (frame_dir, _idx)
|
|
|
|
try:
|
|
os.stat(outfile)
|
|
except FileNotFoundError:
|
|
_frame = _frame.resize((280, 192))
|
|
_frame.save(bmpfile)
|
|
|
|
subprocess.call(
|
|
["/usr/local/bin/bmp2dhr", bmpfile, "hgr", "D9"])
|
|
|
|
os.remove(bmpfile)
|
|
|
|
_frame = np.fromfile(outfile, dtype=np.uint8)
|
|
q.put(_frame)
|
|
|
|
q.put(None)
|
|
|
|
t = threading.Thread(target=worker, daemon=True)
|
|
t.start()
|
|
|
|
while True:
|
|
frame = q.get()
|
|
if frame is None:
|
|
break
|
|
|
|
yield screen.FlatMemoryMap(
|
|
screen_page=1, data=frame).to_memory_map()
|
|
q.task_done()
|
|
|
|
t.join()
|
|
|
|
def encode_frame(
|
|
self, target: screen.MemoryMap
|
|
) -> Iterator[opcodes.Opcode]:
|
|
"""Update to match content of frame within provided budget."""
|
|
|
|
print("Similarity %f" % (self.update_priority.mean()))
|
|
yield from self._index_changes(self.memory_map, target)
|
|
|
|
def _index_changes(
|
|
self,
|
|
source: screen.MemoryMap,
|
|
target: screen.MemoryMap
|
|
) -> Iterator[Tuple[int, int, List[int]]]:
|
|
"""Transform encoded screen to sequence of change tuples."""
|
|
|
|
diff_weights = self._diff_weights(source, target)
|
|
|
|
# Clear any update priority entries that have resolved themselves
|
|
# with new frame
|
|
self.update_priority[diff_weights == 0] = 0
|
|
|
|
# Halve existing weights to increase bias to new diffs.
|
|
# In particular this means that existing updates with diff 1 will
|
|
# become diff 0, i.e. will only be prioritized if they are still
|
|
# diffs in the new frame.
|
|
# self.update_priority >>= 1
|
|
self.update_priority += diff_weights
|
|
|
|
priorities = self._heapify_priorities()
|
|
|
|
content_deltas = {}
|
|
|
|
while priorities:
|
|
_, _, page, offset = heapq.heappop(priorities)
|
|
# Check whether we've already cleared this diff while processing
|
|
# an earlier opcode
|
|
if self.update_priority[page, offset] == 0:
|
|
continue
|
|
|
|
offsets = [offset]
|
|
content = target.page_offset[page, offset]
|
|
|
|
# Clear priority for the offset we're emitting
|
|
self.update_priority[page, offset] = 0
|
|
self.memory_map.page_offset[page, offset] = content
|
|
diff_weights[page, offset] = 0
|
|
|
|
# Make sure we don't emit this offset as a side-effect of some
|
|
# other offset later.
|
|
for cd in content_deltas.values():
|
|
cd[page, offset] = 0
|
|
|
|
# Need to find 3 more offsets to fill this opcode
|
|
for o in self._compute_error(
|
|
page,
|
|
content,
|
|
target,
|
|
diff_weights,
|
|
content_deltas
|
|
):
|
|
offsets.append(o)
|
|
|
|
# Compute new edit distance between new content and target
|
|
# byte, so we can reinsert with this value
|
|
p = edit_distance.edit_weight(
|
|
content, target.page_offset[page, o], o % 2 == 1,
|
|
error=False)
|
|
|
|
# Update priority for the offset we're emitting
|
|
self.update_priority[page, o] = p # 0
|
|
|
|
self.memory_map.page_offset[page, o] = content
|
|
|
|
if p:
|
|
# This content byte introduced an error, so put back on the
|
|
# heap in case we can get back to fixing it exactly
|
|
# during this frame. Otherwise we'll get to it later.
|
|
heapq.heappush(
|
|
priorities, (-p, random.random(), page, offset))
|
|
|
|
# Pad to 4 if we didn't find enough
|
|
for _ in range(len(offsets), 4):
|
|
offsets.append(offsets[0])
|
|
|
|
yield (page + 32, content, offsets)
|
|
|
|
# If we run out of things to do, pad forever
|
|
content = target.page_offset[(0, 0)]
|
|
while True:
|
|
yield (32, content, [0, 0, 0, 0])
|
|
|
|
@staticmethod
|
|
def _diff_weights(
|
|
source: screen.MemoryMap,
|
|
target: screen.MemoryMap
|
|
):
|
|
return edit_distance.screen_edit_distance(
|
|
source.page_offset, target.page_offset)
|
|
|
|
def _heapify_priorities(self) -> List:
|
|
priorities = []
|
|
it = np.nditer(self.update_priority, flags=['multi_index'])
|
|
while not it.finished:
|
|
priority = it[0]
|
|
if not priority:
|
|
it.iternext()
|
|
continue
|
|
|
|
page, offset = it.multi_index
|
|
|
|
# Don't use deterministic order for page, offset
|
|
nonce = random.random()
|
|
priorities.append((-priority, nonce, page, offset))
|
|
it.iternext()
|
|
|
|
heapq.heapify(priorities)
|
|
return priorities
|
|
|
|
@staticmethod
|
|
def _compute_delta(content, target, old):
|
|
"""
|
|
This function is the critical path for the video encoding.
|
|
"""
|
|
return edit_distance.byte_screen_error_distance(content, target) - old
|
|
|
|
_OFFSETS = np.arange(256)
|
|
|
|
def _compute_error(self, page, content, target, old_error, content_deltas):
|
|
offsets = []
|
|
|
|
# TODO: move this up into parent
|
|
delta_screen = content_deltas.get(content)
|
|
if delta_screen is None:
|
|
delta_screen = self._compute_delta(
|
|
content, target.page_offset, old_error)
|
|
content_deltas[content] = delta_screen
|
|
|
|
delta_page = delta_screen[page]
|
|
cond = delta_page < 0
|
|
candidate_offsets = self._OFFSETS[cond]
|
|
priorities = delta_page[cond]
|
|
|
|
l = [
|
|
(priorities[i], random.random(), candidate_offsets[i])
|
|
for i in range(len(candidate_offsets))
|
|
]
|
|
heapq.heapify(l)
|
|
|
|
while l:
|
|
_, _, o = heapq.heappop(l)
|
|
offsets.append(o)
|
|
|
|
# Make sure we don't end up considering this (page, offset) again
|
|
# until the next image frame. Even if a better match comes along,
|
|
# it's probably better to fix up some other byte.
|
|
for cd in content_deltas.values():
|
|
cd[page, o] = 0
|
|
|
|
if len(offsets) == 3:
|
|
break
|
|
|
|
return offsets
|