Revamp frame_grabber - the previous multiprocessing strategy now

causes pickle serialization errors on newer python versions.

WIP - support for using ii-pix for DHGR image conversions instead of bmp2dhr

WIP - support for DHGR mono conversions
This commit is contained in:
kris 2023-01-21 22:08:55 +00:00
parent 1149943d71
commit d131b48661
4 changed files with 174 additions and 89 deletions

View File

@ -28,3 +28,5 @@ threadpoolctl==3.1.0
urllib3==1.26.14
weighted-levenshtein==0.2.1
zipp==3.11.0
colour-science==0.4.1
Cython==0.29.33

View File

@ -1,18 +1,23 @@
"""Extracts sequence of still images from input video stream."""
import os
import queue
import subprocess
import threading
from typing import Iterator
import multiprocessing
from typing import Iterator, Optional
import numpy as np
import skvideo.io
from PIL import Image
from PIL import Image, ImageEnhance
import colour
import screen
from palette import Palette
from video_mode import VideoMode
import convert.palette
import convert.dither_dhr
import convert.dither_pattern
import convert.image
import convert.screen
class FrameGrabber:
@ -50,98 +55,168 @@ class FileFrameGrabber(FrameGrabber):
video_mode.name,
palette.name)
def _palette_arg(self) -> str:
return "P%d" % self.palette.value
def _frame_extractor(self, frame_dir):
for _idx, _frame in enumerate(self._frame_grabber()):
bmpfile = "%s/%08d.bmp" % (frame_dir, _idx)
try:
os.stat(bmpfile)
except FileNotFoundError:
# XXX HGR dimensions
image = _frame.resize((560, 192), resample=Image.LANCZOS)
# XXX gamma?
img_filter = ImageEnhance.Brightness(image)
image = img_filter.enhance(1.5)
image.save(bmpfile)
yield _idx
def frames(self) -> Iterator[screen.MemoryMap]:
"""Encode frame to (D)HGR using bmp2dhr.
"""Encode frame sequence to (D)HGR memory images.
We do the encoding in a background thread to parallelize.
We do the encoding in a worker pool to parallelize.
"""
frame_dir = self._output_dir(
self.filename, self.video_mode, self.palette)
os.makedirs(frame_dir, exist_ok=True)
q = queue.Queue(maxsize=10)
def _hgr_decode(_idx, _frame):
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), resample=Image.LANCZOS)
_frame.save(bmpfile)
subprocess.call([
"/usr/local/bin/bmp2dhr", bmpfile, "hgr",
self._palette_arg(),
"D9" # Buckels dither
])
os.remove(bmpfile)
_main = np.fromfile(outfile, dtype=np.uint8)
return _main, None
def _dhgr_decode(_idx, _frame):
mainfile = "%s/%08d.BIN" % (frame_dir, _idx)
auxfile = "%s/%08d.AUX" % (frame_dir, _idx)
bmpfile = "%s/%08d.bmp" % (frame_dir, _idx)
try:
os.stat(mainfile)
os.stat(auxfile)
except FileNotFoundError:
_frame = _frame.resize((280, 192), resample=Image.LANCZOS)
_frame.save(bmpfile)
subprocess.call([
"/usr/local/bin/bmp2dhr", bmpfile, "dhgr", # "v",
self._palette_arg(),
"A", # Output separate .BIN and .AUX files
"D9" # Buckels dither
])
os.remove(bmpfile)
_main = np.fromfile(mainfile, dtype=np.uint8)
_aux = np.fromfile(auxfile, dtype=np.uint8)
return _main, _aux
def worker():
"""Invoke bmp2dhr to encode input image frames and push to queue."""
decode = (
_dhgr_decode if self.video_mode == VideoMode.DHGR else
_hgr_decode
)
for _idx, _frame in enumerate(self._frame_grabber()):
q.put(decode(_idx, _frame))
q.put((None, None))
t = threading.Thread(target=worker, daemon=True)
t.start()
while True:
main, aux = q.get()
if main is None:
break
global _converter
if self.video_mode == VideoMode.DHGR_MONO:
_converter = DHGRMonoFrameConverter(frame_dir)
elif self.video_mode == VideoMode.DHGR:
# XXX support palette
_converter = DHGRFrameConverter(frame_dir=frame_dir)
elif self.video_mode == VideoMode.HGR:
_converter = HGRFrameConverter(
frame_dir=frame_dir, palette_value=self.palette.value)
pool = multiprocessing.Pool(10)
for main, aux in pool.imap(_converter.convert, self._frame_extractor(
frame_dir), chunksize=1):
main_map = screen.FlatMemoryMap(
screen_page=1, data=main).to_memory_map()
if aux is None:
aux_map = None
else:
aux_map = screen.FlatMemoryMap(
screen_page=1, data=aux).to_memory_map()
yield (main_map, aux_map)
q.task_done()
aux_map = screen.FlatMemoryMap(
screen_page=1, data=aux).to_memory_map() if aux else None
t.join()
yield main_map, aux_map
# Used in worker pool to receive global state
_converter = None # type:Optional[FrameConverter]
class FrameConverter:
def __init__(self, frame_dir):
self.frame_dir = frame_dir
@staticmethod
def convert(idx):
raise NotImplementedError
class HGRFrameConverter(FrameConverter):
def __init__(self, frame_dir, palette_value):
super(HGRFrameConverter, self).__init__(frame_dir)
self.palette_arg = "P%d" % palette_value
@staticmethod
def convert(idx):
outfile = "%s/%08dC.BIN" % (_converter.frame_dir, idx)
bmpfile = "%s/%08d.bmp" % (_converter.frame_dir, idx)
try:
os.stat(outfile)
except FileNotFoundError:
subprocess.call([
"/usr/local/bin/bmp2dhr", bmpfile, "hgr",
_converter.palette_arg,
"D9" # Buckels dither
])
os.remove(bmpfile)
_main = np.fromfile(outfile, dtype=np.uint8)
return _main, None
class DHGRFrameConverter(FrameConverter):
def __init__(self, frame_dir):
super(DHGRFrameConverter, self).__init__(frame_dir)
self.rgb_to_cam16 = np.load(
"transcoder/convert/data/rgb24_to_cam16ucs.npy")
self.dither = convert.dither_pattern.PATTERNS[
convert.dither_pattern.DEFAULT_PATTERN]()
self.convert_palette = convert.palette.PALETTES['ntsc']()
@staticmethod
def convert(idx):
print("Encoding %d" % idx)
bmpfile = "%s/%08d.bmp" % (_converter.frame_dir, idx)
mainfile = "%s/%08d.BIN" % (_converter.frame_dir, idx)
auxfile = "%s/%08d.AUX" % (_converter.frame_dir, idx)
try:
os.stat(mainfile)
os.stat(auxfile)
except FileNotFoundError:
convert_screen = convert.screen.DHGRScreen(
_converter.convert_palette)
# Open and resize source image
image = convert.image.open(bmpfile)
image = np.array(
convert.image.resize(image, convert_screen.X_RES,
convert_screen.Y_RES,
gamma=2.4)).astype(np.float32) / 255
bitmap = convert.dither_dhr.dither_image(
convert_screen, image, _converter.dither, 6,
False, _converter.rgb_to_cam16)
output_screen = convert.screen.DHGRScreen(
_converter.convert_palette)
output_srgb = output_screen.bitmap_to_image_ntsc(bitmap)
out_image = convert.image.resize(
Image.fromarray(output_srgb), convert_screen.X_RES,
convert_screen.Y_RES * 2, srgb_output=True)
outfile = os.path.join(
os.path.splitext(bmpfile)[0] + "-preview.png")
out_image.save(outfile, "PNG")
# out_image.show()
convert_screen.pack(bitmap)
with open(mainfile, "wb") as f:
f.write(bytes(convert_screen.main))
with open(auxfile, "wb") as f:
f.write(bytes(convert_screen.aux))
_main = np.fromfile(mainfile, dtype=np.uint8)
_aux = np.fromfile(auxfile, dtype=np.uint8)
return _main, _aux
class DHGRMonoFrameConverter(FrameConverter):
@staticmethod
def convert(_idx):
bmpfile = "%s/%08d.bmp" % (_converter.frame_dir, _idx)
dhrfile = "%s/%08d.dhr" % (_converter.frame_dir, _idx)
try:
os.stat(dhrfile)
except FileNotFoundError:
subprocess.call([
"python", "convert.py", "dhr_mono", bmpfile, dhrfile
])
# os.remove(bmpfile)
dhr = np.fromfile(dhrfile, dtype=np.uint8)
aux = dhr[:8192]
main = dhr[8192:]
return aux, main
#

View File

@ -84,7 +84,14 @@ class Movie:
aux_memory=aux,
palette=self.palette
)
else:
elif self.video_mode == VideoMode.DHGR_MONO:
# XXX
target_pixelmap = screen.DHGRBitmap(
main_memory=main,
aux_memory=aux,
palette=self.palette
)
else: # VideoMode.HGR
target_pixelmap = screen.HGRBitmap(
main_memory=main,
palette=self.palette
@ -140,7 +147,7 @@ class Movie:
if socket_pos >= 2044:
# 2 op_ack address bytes + 2 payload bytes from ACK must
# terminate 2K stream frame
if self.video_mode == VideoMode.DHGR:
if self.video_mode in {VideoMode.DHGR, VideoMode.DHGR_MONO}:
# Flip-flop between MAIN and AUX banks
self.aux_memory_bank = not self.aux_memory_bank

View File

@ -5,4 +5,5 @@ import enum
class VideoMode(enum.Enum):
HGR = 0 # Hi-Res
DHGR = 1 # Double Hi-Res
DHGR = 1 # Double Hi-Res (Colour)
DHGR_MONO = 2 # Double Hi-Res (Mono)