187 lines
6.5 KiB
Python
187 lines
6.5 KiB
Python
# encoding: utf-8
|
|
"""
|
|
The actual behavior of "Upscaled Crop/Pad" node.
|
|
"""
|
|
|
|
import typing as _t
|
|
|
|
from dataclasses import dataclass as _dataclass
|
|
|
|
from .__dataclass import slots_args as __dataclass_slots_args
|
|
from .funcs import round_pos_int as _round_pos_int, _need_post_resize, _show_text_on_node
|
|
from .enums import *
|
|
from .return_tuples import *
|
|
|
|
|
|
@_dataclass(**__dataclass_slots_args)
|
|
class _CropPadInput:
|
|
"""
|
|
For the ease of passing between the functions, the entire set of inputs for "Upscaled Crop/Pad" node
|
|
is internally turned into a dataclass.
|
|
"""
|
|
upscale: float
|
|
|
|
init_w: int
|
|
init_h: int
|
|
hd_w: int
|
|
hd_h: int
|
|
|
|
strategy: str
|
|
align_x: float
|
|
align_y: float
|
|
|
|
|
|
def _upscaled_crop_xy_offset(extra_w: int, extra_h: int, align_x: float, align_y: float):
|
|
crop_x_origin = _round_pos_int(align_x * extra_w)
|
|
crop_y_origin = _round_pos_int((1.0 - align_y) * extra_h)
|
|
return crop_x_origin, crop_y_origin
|
|
|
|
|
|
def _upscaled_crop(_in: _CropPadInput, real_upscale: float) -> ResultUpscaledCropPad:
|
|
raw_upscaled_w = _round_pos_int(real_upscale * _in.init_w)
|
|
raw_upscaled_h = _round_pos_int(real_upscale * _in.init_h)
|
|
extra_w = max(raw_upscaled_w - _in.hd_w, 0)
|
|
extra_h = max(raw_upscaled_h - _in.hd_h, 0)
|
|
crop_x_offset, crop_y_offset = _upscaled_crop_xy_offset(extra_w, extra_h, _in.align_x, _in.align_y)
|
|
return ResultUpscaledCropPad(
|
|
real_upscale,
|
|
True, _in.hd_w, _in.hd_h, crop_x_offset, crop_y_offset,
|
|
False, 0, 0, 0, 0,
|
|
)
|
|
|
|
|
|
def _upscaled_pad_side_values(extra_w: int, extra_h: int, align_x: float, align_y: float):
|
|
pad_left = _round_pos_int(align_x * extra_w)
|
|
pad_right = extra_w - pad_left
|
|
pad_bottom = _round_pos_int(align_y * extra_h)
|
|
pad_top = extra_h - pad_bottom
|
|
return pad_left, pad_top, pad_right, pad_bottom
|
|
|
|
|
|
def _upscaled_pad(_in: _CropPadInput, real_upscale: float) -> ResultUpscaledCropPad:
|
|
raw_upscaled_w = _round_pos_int(real_upscale * _in.init_w)
|
|
raw_upscaled_h = _round_pos_int(real_upscale * _in.init_h)
|
|
pad_w = max(_in.hd_w - raw_upscaled_w, 0)
|
|
pad_h = max(_in.hd_h - raw_upscaled_h, 0)
|
|
return ResultUpscaledCropPad(
|
|
real_upscale,
|
|
False, raw_upscaled_w, raw_upscaled_h, 0, 0,
|
|
True, *_upscaled_pad_side_values(pad_w, pad_h, _in.align_x, _in.align_y),
|
|
)
|
|
|
|
|
|
def _upscaled_crop_delta_pixels(_in: _CropPadInput, result: ResultUpscaledCropPad) -> int:
|
|
"""Calculate the area of total cropped patch for crop-only mode."""
|
|
raw_upscaled_w = _round_pos_int(result.upscale * _in.init_w)
|
|
raw_upscaled_h = _round_pos_int(result.upscale * _in.init_h)
|
|
cropped_w = min(result.crop_width, raw_upscaled_w)
|
|
cropped_h = min(result.crop_height, raw_upscaled_h)
|
|
return abs(raw_upscaled_w * raw_upscaled_h - cropped_w * cropped_h)
|
|
|
|
|
|
def _upscaled_pad_delta_pixels(_in: _CropPadInput, result: ResultUpscaledCropPad) -> int:
|
|
"""Calculate the area of total out-painted patch for pad-only mode."""
|
|
raw_upscaled_w = _round_pos_int(result.upscale * _in.init_w)
|
|
raw_upscaled_h = _round_pos_int(result.upscale * _in.init_h)
|
|
padded_w = raw_upscaled_w + max(result.pad_left, 0) + max(result.pad_right, 0)
|
|
padded_h = raw_upscaled_h + max(result.pad_top, 0) + max(result.pad_bottom, 0)
|
|
return abs(padded_w * padded_h - raw_upscaled_w * raw_upscaled_h)
|
|
|
|
|
|
def _upscaled_crop_pad(_in: _CropPadInput) -> ResultUpscaledCropPad:
|
|
"""The actual function for "Upscaled Crop/Pad" node - extracted to wrap it with displaying the message."""
|
|
needs_resize, real_upscale_avg, real_upscale_x, real_upscale_y = _need_post_resize(_in.init_w, _in.init_h, _in.hd_w, _in.hd_h)
|
|
if not needs_resize:
|
|
return ResultUpscaledCropPad(
|
|
real_upscale_avg,
|
|
False, _in.hd_w, _in.hd_h, 0, 0,
|
|
False, 0, 0, 0, 0,
|
|
)
|
|
|
|
_in.align_x = min(max(_in.align_x, 0.0), 1.0)
|
|
_in.align_y = min(max(_in.align_y, 0.0), 1.0)
|
|
|
|
if _in.strategy == UpscaledCropPadStrategy.CROP:
|
|
return _upscaled_crop(_in, max(real_upscale_x, real_upscale_y))
|
|
|
|
if _in.strategy == UpscaledCropPadStrategy.PAD:
|
|
return _upscaled_pad(_in, min(real_upscale_x, real_upscale_y))
|
|
|
|
if _in.strategy == UpscaledCropPadStrategy.NEAREST:
|
|
crop_result = _upscaled_crop(_in, max(real_upscale_x, real_upscale_y))
|
|
pad_result = _upscaled_pad(_in, min(real_upscale_x, real_upscale_y))
|
|
return (
|
|
pad_result
|
|
if _upscaled_pad_delta_pixels(_in, pad_result) < _upscaled_crop_delta_pixels(_in, crop_result)
|
|
else crop_result
|
|
)
|
|
|
|
assert _in.strategy == UpscaledCropPadStrategy.EXACT_UPSCALE
|
|
|
|
upscale = _in.upscale
|
|
raw_upscaled_w = _round_pos_int(upscale * _in.init_w)
|
|
raw_upscaled_h = _round_pos_int(upscale * _in.init_h)
|
|
# We're in the most complex case. These ^ can be both below and above the target up-res.
|
|
# So, we need to follow the whole upscale->crop->pad chain: we'll just get zeroes where there will be nothing to do.
|
|
|
|
cropped_w = min(_in.hd_w, raw_upscaled_w)
|
|
cropped_h = min(_in.hd_h, raw_upscaled_h)
|
|
crop_extra_w = raw_upscaled_w - cropped_w
|
|
crop_extra_h = raw_upscaled_h - cropped_h
|
|
crop_x_offset, crop_y_offset = _upscaled_crop_xy_offset(crop_extra_w, crop_extra_h, _in.align_x, _in.align_y)
|
|
|
|
pad_extra_w = _in.hd_w - cropped_w
|
|
pad_extra_h = _in.hd_h - cropped_h
|
|
pad_side_values = _upscaled_pad_side_values(pad_extra_w, pad_extra_h, _in.align_x, _in.align_y)
|
|
return ResultUpscaledCropPad(
|
|
upscale,
|
|
crop_extra_w > 0 or crop_extra_h > 0, cropped_w, cropped_h, crop_x_offset, crop_y_offset,
|
|
pad_extra_w > 0 or pad_extra_h > 0, *pad_side_values,
|
|
)
|
|
|
|
|
|
def upscaled_crop_pad(
|
|
upscale: float,
|
|
init_w: int, init_h: int, hd_w: int, hd_h: int,
|
|
strategy: str,
|
|
align_x: float, align_y: float,
|
|
show: bool = True,
|
|
unique_id: str = None
|
|
) -> ResultUpscaledCropPad:
|
|
"""
|
|
Detect, whether some cropping/out-painting needs to be done after upscaling the original image.
|
|
And if so - what are the values for them.
|
|
"""
|
|
# noinspection PyArgumentList
|
|
result = _upscaled_crop_pad(_CropPadInput(
|
|
upscale,
|
|
init_w, init_h, hd_w, hd_h,
|
|
strategy,
|
|
align_x, align_y,
|
|
))
|
|
|
|
if not unique_id:
|
|
return result
|
|
|
|
if not show:
|
|
_show_text_on_node(None, unique_id)
|
|
return result
|
|
|
|
if not(result.do_crop or result.do_padding):
|
|
_show_text_on_node(f"x{result.upscale:.3f}✅", unique_id)
|
|
return result
|
|
|
|
# Either crop, or pad, or both
|
|
text_parts: _t.List[str] = [f"x{result.upscale:.3f}", ]
|
|
|
|
if result.do_crop:
|
|
crop_r = hd_w - (result.crop_x_origin + result.crop_width)
|
|
crop_b = hd_h - (result.crop_y_origin + result.crop_height)
|
|
text_parts.append(f"Crop:\nL {result.crop_x_origin}, R {crop_r}, B {crop_b}, T {result.crop_y_origin}")
|
|
|
|
if result.do_padding:
|
|
text_parts.append(f"Padding:\nL {result.pad_left}, R {result.pad_right}, B {result.pad_bottom}, T {result.pad_top}")
|
|
|
|
_show_text_on_node('\n\n'.join(text_parts), unique_id)
|
|
return result
|