FrameForge/image_generator/__init__.py
2024-01-24 12:11:06 -05:00

586 lines
21 KiB
Python

import re
import shutil
from random import choice
from pathlib import Path
from numpy import linspace
from unidecode import unidecode
import awsmfunc
import vapoursynth as vs
from image_generator.exceptions import ImageGeneratorError
from image_generator.utils import get_working_dir, hex_to_bgr
class GenerateImages:
def __init__(
self,
source_file: Path,
encode_file: Path,
image_dir: Path,
indexer: str,
index_directory: None | str,
sub_size: int,
left_crop: int,
right_crop: int,
top_crop: int,
bottom_crop: int,
adv_resize_left: float,
adv_resize_right: float,
adv_resize_top: float,
adv_resize_bottom: float,
tone_map: bool,
re_sync: str,
comparison_count: int,
subtitle_color: str,
release_sub_title: str | None,
):
self.source_file = source_file
self.source_node = None
self.reference_source_file = None
self.encode_file = encode_file
self.encode_node = None
self.image_dir = image_dir
self.indexer = indexer
self.index_dir = index_directory
self.sub_size = sub_size
self.left_crop = left_crop
self.right_crop = right_crop
self.top_crop = top_crop
self.bottom_crop = bottom_crop
self.adv_resize_left = adv_resize_left
self.adv_resize_right = adv_resize_right
self.adv_resize_top = adv_resize_top
self.adv_resize_bottom = adv_resize_bottom
self.tone_map = tone_map
self.re_sync = re_sync
self.comparison_count = comparison_count
self.subtitle_color = subtitle_color
self.release_sub_title = release_sub_title
self.core = vs.core
self.load_plugins()
def process_images(self):
if self.indexer == "lsmash":
self.index_lsmash()
elif self.indexer == "ffms2":
self.index_ffms2()
num_source_frames = len(self.source_node)
num_encode_frames = len(self.encode_node)
# Format: Name, Fontname, Font-size, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic,
# Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL,
# MarginR, MarginV,
# bgr color
color = "&H000ac7f5"
if self.subtitle_color:
color = hex_to_bgr(self.subtitle_color)
selected_sub_style = (
f"Segoe UI,{self.sub_size},{color},&H00000000,&H00000000,&H00000000,"
"1,0,0,0,100,100,0,0,1,1,0,7,5,0,0,1"
)
self.check_de_interlaced(num_source_frames, num_encode_frames)
b_frames = self.get_b_frames(num_source_frames)
screenshot_comparison_dir, screenshot_sync_dir = self.generate_folders()
self.handle_crop()
self.handle_resize()
self.handle_hdr()
vs_source_info, vs_encode_info = self.handle_subtitles(selected_sub_style)
img_job = self.generate_screens(
b_frames,
vs_source_info,
vs_encode_info,
screenshot_comparison_dir,
screenshot_sync_dir,
selected_sub_style,
)
return img_job
@staticmethod
def screen_gen_callback(sg_call_back):
print(
str(sg_call_back).replace("ScreenGen: ", "").replace("\n", "").strip(),
flush=True,
)
def generate_screens(
self,
b_frames,
vs_source_info,
vs_encode_info,
screenshot_comparison_dir,
screenshot_sync_dir,
selected_sub_style,
):
print("\nGenerating screenshots, please wait", flush=True)
# handle re_sync if needed
sync_frames = []
if self.re_sync:
get_sync_digits = re.search(r"(\d+)", self.re_sync)
sync_digits = int(get_sync_digits.group(1)) if get_sync_digits else 0
for x_frames in b_frames:
if self.re_sync.startswith("-"):
sync_frames.append(int(x_frames) - sync_digits)
else:
sync_frames.append(int(x_frames) + sync_digits)
else:
sync_frames = b_frames
# generate source images
awsmfunc.ScreenGen(
vs_source_info,
frame_numbers=sync_frames,
fpng_compression=1,
folder=screenshot_comparison_dir,
suffix="a_source__%d",
callback=self.screen_gen_callback,
)
# generate encode images
awsmfunc.ScreenGen(
vs_encode_info,
frame_numbers=b_frames,
fpng_compression=1,
folder=screenshot_comparison_dir,
suffix="b_encode__%d",
callback=self.screen_gen_callback,
)
# generate some sync frames
print("\nGenerating a few sync frames", flush=True)
# select two frames randomly from list
sync_1 = choice(b_frames)
remove_sync1 = b_frames.copy()
remove_sync1.remove(sync_1)
sync_2 = choice(remove_sync1)
# reference subs
vs_source_ref_info = self.core.sub.Subtitle(
clip=self.source_node, text="Sync", style=selected_sub_style
)
vs_encode_ref_info = self.core.sub.Subtitle(
clip=self.encode_node, text="Reference", style=selected_sub_style
)
# generate screens for the two reference frames
awsmfunc.ScreenGen(
vs_encode_ref_info,
frame_numbers=[sync_1, sync_2],
fpng_compression=1,
folder=screenshot_sync_dir,
suffix="b_encode__%d",
callback=self.screen_gen_callback,
)
# generate 10 source frames around those reference frames
awsmfunc.ScreenGen(
vs_source_ref_info,
frame_numbers=[
sync_1 - 4,
sync_1 - 3,
sync_1 - 2,
sync_1 - 1,
sync_1,
sync_1 + 1,
sync_1 + 2,
sync_1 + 3,
sync_1 + 4,
],
fpng_compression=1,
folder=Path(Path(screenshot_sync_dir) / "sync1"),
suffix="a_source__%d",
callback=self.screen_gen_callback,
)
awsmfunc.ScreenGen(
vs_source_ref_info,
frame_numbers=[
sync_2 - 4,
sync_2 - 3,
sync_2 - 2,
sync_2 - 1,
sync_2,
sync_2 + 1,
sync_2 + 2,
sync_2 + 3,
sync_2 + 4,
],
fpng_compression=1,
folder=Path(Path(screenshot_sync_dir) / "sync2"),
suffix="a_source__%d",
callback=self.screen_gen_callback,
)
print("Screen generation completed", flush=True)
return str(screenshot_comparison_dir)
def handle_subtitles(self, selected_sub_style):
vs_source_info = self.core.sub.Subtitle(
clip=self.source_node, text="Source", style=selected_sub_style
)
vs_encode_info = awsmfunc.FrameInfo(
clip=self.encode_node,
title=self.release_sub_title if self.release_sub_title else "",
style=selected_sub_style,
)
return vs_source_info, vs_encode_info
def handle_hdr(self):
if self.tone_map:
self.source_node = awsmfunc.DynamicTonemap(
clip=self.source_node, libplacebo=False
)
self.encode_node = awsmfunc.DynamicTonemap(
clip=self.encode_node,
reference=self.reference_source_file,
libplacebo=False,
)
def handle_resize(self):
if (
self.source_node.width != self.encode_node.width
and self.source_node.height != self.encode_node.height
or any(
[
self.adv_resize_left,
self.adv_resize_right,
self.adv_resize_top,
self.adv_resize_bottom,
]
)
):
# advanced resize offset vars
advanced_resize_left = self.adv_resize_left if self.adv_resize_left else 0
advanced_resize_top = self.adv_resize_top if self.adv_resize_top else 0
advanced_resize_width = (
self.adv_resize_right if self.adv_resize_right else 0
)
advanced_resize_height = (
self.adv_resize_bottom if self.adv_resize_bottom else 0
)
# resize source to match encode for screenshots
self.source_node = self.core.resize.Spline36(
self.source_node,
width=int(self.encode_node.width),
height=int(self.encode_node.height),
src_left=advanced_resize_left,
src_top=advanced_resize_top,
src_width=float(
self.source_node.width
- (advanced_resize_left + advanced_resize_width)
),
src_height=float(
self.source_node.height
- (advanced_resize_top + advanced_resize_height)
),
dither_type="error_diffusion",
)
def handle_crop(self):
if any([self.left_crop, self.right_crop, self.top_crop, self.bottom_crop]):
self.source_node = self.core.std.Crop(
self.source_node,
left=self.left_crop if self.left_crop else 0,
right=self.right_crop if self.right_crop else 0,
top=self.top_crop if self.top_crop else 0,
bottom=self.bottom_crop if self.bottom_crop else 0,
)
def generate_folders(self):
print("\nCreating folders for images", flush=True)
if self.image_dir:
image_output_dir = Path(self.image_dir)
else:
image_output_dir = Path(
Path(self.encode_file).parent / f"{Path(self.encode_file).stem}_images"
)
# remove any accent characters from path
image_output_dir = Path(unidecode(str(image_output_dir)))
# check if temp image dir exists, if so delete it!
if image_output_dir.exists():
shutil.rmtree(image_output_dir, ignore_errors=True)
# create main image dir
image_output_dir.mkdir(exist_ok=True, parents=True)
# create comparison image directory and define it as variable
Path(Path(image_output_dir) / "img_comparison").mkdir(exist_ok=True)
screenshot_comparison_dir = str(Path(Path(image_output_dir) / "img_comparison"))
# create selected image directory and define it as variable
Path(Path(image_output_dir) / "img_selected").mkdir(exist_ok=True)
# create sync image directory and define it as variable
Path(Path(image_output_dir) / "img_sync").mkdir(exist_ok=True)
screenshot_sync_dir = str(Path(Path(image_output_dir) / "img_sync"))
# create sub directories
Path(Path(image_output_dir) / "img_sync/sync1").mkdir(exist_ok=True)
Path(Path(image_output_dir) / "img_sync/sync2").mkdir(exist_ok=True)
print("Folder creation completed", flush=True)
return screenshot_comparison_dir, screenshot_sync_dir
def get_b_frames(self, num_source_frames):
print(
f"\nGenerating {self.comparison_count} 'B' frames for " "comparison images",
flush=True,
)
b_frames = list(
linspace(
int(num_source_frames * 0.15),
int(num_source_frames * 0.75),
int(self.comparison_count),
).astype(int)
)
try:
for i, frame in enumerate(b_frames):
while (
self.encode_node.get_frame(frame).props["_PictType"].decode() != "B"
):
frame += 1
b_frames[i] = frame
except ValueError:
raise ImageGeneratorError(
"Error! Your encode file is likely an incomplete or corrupted encode"
)
print(f"Finished generating {self.comparison_count} 'B' frames", flush=True)
return b_frames
def check_de_interlaced(self, num_source_frames, num_encode_frames):
print("\nChecking if encode has been de-interlaced", flush=True)
try:
source_fps = float(self.source_node.fps)
encode_fps = float(self.encode_node.fps)
if source_fps != encode_fps:
if num_source_frames == num_encode_frames:
self.source_node = self.core.std.AssumeFPS(
self.source_node,
fpsnum=self.encode_node.fps.numerator,
fpsden=self.encode_node.fps.denominator,
)
print(
"Adjusting source fps to match the encode using AssumeFPS() on the source",
flush=True,
)
else:
even_frames_for = ""
if num_source_frames != num_encode_frames:
file_differences = float(num_source_frames / num_encode_frames)
if file_differences > 1.01:
even_frames_for = "source"
self.source_node = self.core.std.SelectEvery(
self.source_node, cycle=2, offsets=0
)
elif file_differences < 0.99:
even_frames_for = "encode"
self.encode_node = self.core.std.SelectEvery(
self.encode_node, cycle=2, offsets=0
)
print(
f"Source: FPS={source_fps} Frames={num_source_frames}\n"
f"Encode: FPS={encode_fps} Frames={num_encode_frames}\n"
"Source vs Encode appears to be different, we're going to assume encode has been"
f" de-interlaced, automatically generating even frames for {even_frames_for}",
flush=True,
)
else:
print("No de-interlacing detected", flush=True)
except ValueError:
print(
"There was an error while detecting source or encode fps, attempting to continue",
flush=True,
)
def index_lsmash(self):
print("Indexing source", flush=True)
# index source file
# if index is found in the StaxRip temp working directory, attempt to use it
if (
Path(str(Path(self.source_file).with_suffix("")) + "_temp/").is_dir()
and Path(
str(Path(self.source_file).with_suffix("")) + "_temp/temp.lwi"
).is_file()
):
print("Index found in StaxRip temp, attempting to use", flush=True)
# define cache path
lwi_cache_path = Path(
str(Path(self.source_file).with_suffix("")) + "_temp/temp.lwi"
)
# try to use index on source file with the cache path
try:
self.source_node = self.core.lsmas.LWLibavSource(
source=self.source_file, cachefile=lwi_cache_path
)
self.reference_source_file = self.core.lsmas.LWLibavSource(
source=self.source_file, cachefile=lwi_cache_path
)
print("Using existing index", flush=True)
# if index cannot be used
except vs.Error:
print("L-Smash version miss-match, indexing source again", flush=True)
# index source file
self.source_node = self.core.lsmas.LWLibavSource(self.source_file)
self.reference_source_file = self.core.lsmas.LWLibavSource(
self.source_file
)
# if no existing index is found index source file
else:
cache_path = Path(Path(self.source_file).with_suffix(".lwi"))
try:
# create index
self.source_node = self.core.lsmas.LWLibavSource(
self.source_file, cachefile=cache_path
)
self.reference_source_file = self.core.lsmas.LWLibavSource(
self.source_file, cachefile=cache_path
)
except vs.Error:
# delete index
Path(self.source_file).with_suffix(".lwi").unlink(missing_ok=True)
# create index
self.source_node = self.core.lsmas.LWLibavSource(
self.source_file, cachefile=cache_path
)
self.reference_source_file = self.core.lsmas.LWLibavSource(
self.source_file, cachefile=cache_path
)
print("Source index completed\n\nIndexing encode", flush=True)
# define a path for encode index to go
if self.index_dir:
index_base_path = Path(self.index_dir) / Path(self.encode_file).name
cache_path_enc = index_base_path.with_suffix(".lwi")
else:
cache_path_enc = Path(Path(self.encode_file).with_suffix(".lwi"))
try:
# create index
self.encode_node = self.core.lsmas.LWLibavSource(
self.encode_file, cachefile=cache_path_enc
)
except vs.Error:
# delete index
cache_path_enc.unlink(missing_ok=True)
# create index
self.encode_node = self.core.lsmas.LWLibavSource(
self.encode_file, cachefile=cache_path_enc
)
print("Encode index completed", flush=True)
def index_ffms2(self):
print("Indexing source", flush=True)
# index source file
# if index is found in the StaxRip temp working directory, attempt to use it
if (
Path(str(Path(self.source_file).with_suffix("")) + "_temp/").is_dir()
and Path(
str(Path(self.source_file).with_suffix("")) + "_temp/temp.ffindex"
).is_file()
):
print("Index found in StaxRip temp, attempting to use", flush=True)
# define cache path
ffindex_cache_path = Path(
str(Path(self.source_file).with_suffix("")) + "_temp/temp.ffindex"
)
# try to use index on source file with the cache path
try:
self.source_node = self.core.ffms2.Source(
source=self.source_file, cachefile=ffindex_cache_path
)
self.reference_source_file = self.core.ffms2.Source(
source=self.source_file, cachefile=ffindex_cache_path
)
print("Using existing index", flush=True)
# if index cannot be used
except vs.Error:
print("FFMS2 version miss-match, indexing source again", flush=True)
# index source file
self.source_node = self.core.ffms2.Source(self.source_file)
self.reference_source_file = self.core.ffms2.Source(self.source_file)
# if no existing index is found index source file
else:
try:
# create index
print(
"FFMS2 library doesn't allow progress, please wait while the index is completed",
flush=True,
)
self.source_node = self.core.ffms2.Source(self.source_file)
self.reference_source_file = self.core.ffms2.Source(self.source_file)
except vs.Error:
# delete index
Path(self.source_file).with_suffix(".ffindex").unlink(missing_ok=True)
# create index
print(
"FFMS2 library doesn't allow progress, please wait while the index is completed",
flush=True,
)
self.source_node = self.core.ffms2.Source(self.source_file)
self.reference_source_file = self.core.ffms2.Source(self.source_file)
print("Source index completed\n\nIndexing encode", flush=True)
# define a path for encode index to go
if self.index_dir:
index_base_path = Path(self.index_dir) / Path(self.encode_file).name
cache_path_enc = Path(str(index_base_path) + ".ffindex")
else:
cache_path_enc = Path(self.encode_file + ".ffindex")
try:
self.encode_node = self.core.ffms2.Source(
self.encode_file, cachefile=cache_path_enc
)
except vs.Error:
cache_path_enc.unlink(missing_ok=True)
self.encode_node = self.core.ffms2.Source(
self.encode_file, cachefile=cache_path_enc
)
print("Encode index completed", flush=True)
def load_plugins(self):
plugin_path = get_working_dir() / "img_plugins"
if not plugin_path.is_dir() and not plugin_path.exists():
raise ImageGeneratorError("Can not detect plugin directory")
else:
for plugin in plugin_path.glob("*.dll"):
self.core.std.LoadPlugin(Path(plugin).resolve())