|
import argparse |
|
import sys |
|
|
|
from PIL import Image |
|
from typing import List, Optional, Tuple |
|
|
|
|
|
Pos = Tuple[int, int] |
|
Dim = Tuple[int, int] |
|
|
|
|
|
class Box: |
|
def __init__(self, min: Pos, max: Pos) -> None: |
|
self._min = min |
|
self._max = max |
|
|
|
|
|
def min(self) -> Tuple[int, int]: |
|
return self._min |
|
|
|
|
|
def max(self) -> Tuple[int, int]: |
|
return self._max |
|
|
|
def width(self) -> int: |
|
return self._max[0] - self._min[0] + 1 |
|
|
|
def height(self) -> int: |
|
return self._max[1] - self._min[1] + 1 |
|
|
|
def dimensions(self) -> Tuple[int, int]: |
|
return (self.width(), self.height()) |
|
|
|
|
|
def as_tuple(self) -> Tuple[int, int, int, int]: |
|
return (self._min[0], self._min[1], self._max[0], self._max[1]) |
|
|
|
|
|
class DownBox(Box): |
|
def __init__(self, min: Pos, max: Pos, down_pos: Pos) -> None: |
|
super().__init__(min, max) |
|
self._down_pos = down_pos |
|
|
|
def down_pos(self) -> Tuple[int, int]: |
|
return self._down_pos |
|
|
|
|
|
class ExtractedBoxes: |
|
def __init__(self, boxes: List[DownBox]) -> None: |
|
self._boxes = boxes |
|
|
|
def boxes(self) -> List[DownBox]: |
|
return self._boxes |
|
|
|
def down_dimensions(self) -> Dim: |
|
if len(self._boxes) == 0: |
|
return (0, 0) |
|
back = self._boxes[-1] |
|
down = back.down_pos() |
|
return (down[0] + 1, down[1] + 1) |
|
|
|
def full_dimensions(self) -> Dim: |
|
if len(self._boxes) == 0: |
|
return (0, 0) |
|
back = self._boxes[-1] |
|
max = back.max() |
|
return (max[0] + 1, max[1] + 1) |
|
|
|
def to_colored_checkers(self, *, full=True) -> Image.Image: |
|
if full: |
|
width, height = self.full_dimensions() |
|
else: |
|
width, height = self.down_dimensions() |
|
if width == 0 or height == 0: |
|
return Image.new("RGB", (0, 0)) |
|
image = Image.new("RGB", (width, height)) |
|
colors = [ |
|
(255, 255, 255), |
|
(0, 0, 0), |
|
(255, 0, 0), |
|
(255, 127, 0), |
|
(255, 255, 0), |
|
(0, 255, 0), |
|
(0, 0, 255), |
|
(75, 0, 130), |
|
(148, 0, 211), |
|
(255, 0, 255), |
|
] |
|
colorsMax = len(colors) |
|
currColor = 0 |
|
for box in self._boxes: |
|
color = colors[currColor] |
|
currColor = (currColor + 1) % colorsMax |
|
if full: |
|
dim = box.dimensions() |
|
pos = box.min() |
|
else: |
|
dim = (1, 1) |
|
pos = box.down_pos() |
|
subImage = Image.new("RGB", dim, color) |
|
image.paste(subImage, pos) |
|
return image |
|
|
|
|
|
def average_box_dimensions(boxes: List[DownBox]) -> Dim: |
|
assert len(boxes) > 0 |
|
if len(boxes) == 1: |
|
return boxes[0].dimensions() |
|
if len(boxes) <= 16: |
|
|
|
width = 0 |
|
height = 0 |
|
for box in boxes: |
|
width += box.width() |
|
height += box.height() |
|
return (width // len(boxes), height // len(boxes)) |
|
|
|
widths = [box.width() for box in boxes] |
|
heights = [box.height() for box in boxes] |
|
widths.sort() |
|
heights.sort() |
|
return (widths[len(widths) // 2], heights[len(heights) // 2]) |
|
|
|
|
|
def get_trimmed(boxes: List[DownBox]) -> Tuple[Box, Box]: |
|
avg = average_box_dimensions(boxes) |
|
|
|
outlier_dist = 1 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def is_outlier(box: DownBox) -> bool: |
|
dim = box.dimensions() |
|
if abs(dim[0] - avg[0]) > outlier_dist: |
|
return True |
|
if abs(dim[1] - avg[1]) > outlier_dist: |
|
return True |
|
return False |
|
|
|
assert len(boxes) > 0 |
|
front = boxes[0] |
|
back = boxes[-1] |
|
|
|
min_out = (0, 0) |
|
max_out = back.max() |
|
min_down = (0, 0) |
|
max_down = back.down_pos() |
|
if is_outlier(front): |
|
for i in range(1, len(boxes)): |
|
if not is_outlier(boxes[i]): |
|
min_out = boxes[i].min() |
|
min_down = boxes[i].down_pos() |
|
break |
|
if is_outlier(back): |
|
for i in range(len(boxes) - 2, -1, -1): |
|
if not is_outlier(boxes[i]): |
|
max_out = boxes[i].max() |
|
max_down = boxes[i].down_pos() |
|
break |
|
box_out = Box(min_out, max_out) |
|
box_down = Box(min_down, max_down) |
|
return (box_out, box_down) |
|
|
|
|
|
def calc_face_box(control_image: Image.Image, min_pos: Pos) -> Box: |
|
min_pixel = control_image.getpixel(min_pos) |
|
width, height = control_image.size |
|
x = 0 |
|
while min_pos[0] + x < width: |
|
if control_image.getpixel((min_pos[0] + x, min_pos[1])) != min_pixel: |
|
break |
|
x += 1 |
|
y = 0 |
|
while min_pos[1] + y < height: |
|
if control_image.getpixel((min_pos[0], min_pos[1] + y)) != min_pixel: |
|
break |
|
y += 1 |
|
x -= 1 |
|
y -= 1 |
|
assert x > 0 |
|
assert y > 0 |
|
return Box(min_pos, (x + min_pos[0], y + min_pos[1])) |
|
|
|
|
|
def extract_boxes(control_image: Image.Image) -> ExtractedBoxes: |
|
width, height = control_image.size |
|
assert width > 0 |
|
assert height > 0 |
|
|
|
boxes: List[DownBox] = [] |
|
x = 0 |
|
y = 0 |
|
down_x = 0 |
|
down_y = 0 |
|
|
|
while y < height: |
|
while x < width: |
|
min_pos = (x, y) |
|
box = calc_face_box(control_image, min_pos) |
|
boxes.append(DownBox(box.min(), box.max(), (down_x, down_y))) |
|
x += box.width() |
|
down_x += 1 |
|
assert x == width |
|
box = boxes[-1] |
|
x = 0 |
|
y += box.height() |
|
down_x = 0 |
|
down_y += 1 |
|
assert y == height |
|
|
|
return ExtractedBoxes(boxes) |
|
|
|
|
|
def downsample_one(input_image: Image.Image, box: Box, sample_radius: Optional[int], downsampler: Image.Resampling) -> Tuple[int, int, int]: |
|
region = input_image.crop(box.as_tuple()) |
|
|
|
box_width = box.width() |
|
box_height = box.height() |
|
box_center_x = box.min()[0] + box_width // 2 |
|
box_center_y = box.min()[1] + box_height // 2 |
|
|
|
if sample_radius is not None: |
|
radius_x = min(sample_radius, box_width // 2) |
|
radius_y = min(sample_radius, box_height // 2) |
|
else: |
|
radius_x = box_width // 2 |
|
radius_y = box_height // 2 |
|
|
|
cropped_region = region.crop(( |
|
max(0, box_center_x - radius_x - box.min()[0]), |
|
max(0, box_center_y - radius_y - box.min()[1]), |
|
min(box_width, box_center_x + radius_x - box.min()[0]), |
|
min(box_height, box_center_y + radius_y - box.min()[1]) |
|
)) |
|
assert cropped_region.size[0] >= radius_x and cropped_region.size[1] >= radius_y |
|
sampled = cropped_region.resize((1, 1), downsampler) |
|
|
|
rgb_value = sampled.getpixel((0, 0)) |
|
assert isinstance(rgb_value, tuple) and len(rgb_value) == 3 |
|
return rgb_value |
|
|
|
|
|
class ImageRef: |
|
def __init__(self, ref: Image.Image) -> None: |
|
self.ref = ref |
|
|
|
|
|
def downsample_all(*, input_image: Image.Image, output_image: Optional[ImageRef], down_image: Optional[ImageRef], boxes: List[DownBox], sample_radius: Optional[int], downsampler: Image.Resampling, trim_cropped_edges: bool) -> None: |
|
assert output_image or down_image |
|
for box in boxes: |
|
rgb_value = downsample_one(input_image, box, sample_radius, downsampler) |
|
solid_color_image = Image.new("RGB", box.dimensions(), rgb_value) |
|
if output_image: |
|
output_image.ref.paste(solid_color_image, box.min()) |
|
if down_image: |
|
down_image.ref.paste(solid_color_image, box.down_pos()) |
|
if trim_cropped_edges: |
|
o, d = get_trimmed(boxes) |
|
if output_image: |
|
output_image.ref = output_image.ref.crop(o.as_tuple()) |
|
if down_image: |
|
down_image.ref = down_image.ref.crop(d.as_tuple()) |
|
|
|
|
|
def str2bool(value) -> bool: |
|
if isinstance(value, bool): |
|
return value |
|
if value.lower() in ("true", "1"): |
|
return True |
|
elif value.lower() in ("false", "0"): |
|
return False |
|
else: |
|
raise argparse.ArgumentTypeError("Boolean value expected.") |
|
|
|
|
|
def controlled_downscale(*, control_path: str, input_path: str, output_downscaled_path: Optional[str], output_quantized_path: Optional[str], sample_radius: Optional[int], downsampler: Image.Resampling, trim_cropped_edges: bool, output_colorized_full_path: Optional[str], output_colorized_down_path: Optional[str]) -> None: |
|
""" |
|
Downsample and rescale an image. |
|
|
|
:param control_path: Path to the control image. |
|
:param input_path: Path to the input image. |
|
:param output_downscaled_path: Path to save the output downscaled image. |
|
:param output_quantized_path: Path to save the output quantized image (downscaled and then upscaled to the original size). |
|
:param sample_radius: Radius for sampling (Manhattan distance). |
|
:param downsampler: Downsampler to use. |
|
:param trim_cropped_edges: Drop mapped checker grid elements that are cropped in the control image. |
|
:param output_colorized_full_path: Colorize the full checker image to debug the checker parsing. |
|
:param output_colorized_down_path: Colorize the downscaled checker image to debug the checker parsing. |
|
""" |
|
if not output_downscaled_path and not output_quantized_path: |
|
raise ValueError("At least one of output_up and output_down must be specified.") |
|
|
|
control_image = Image.open(control_path).convert("1") |
|
input_image = Image.open(input_path) |
|
if control_image.size != input_image.size: |
|
raise ValueError("Control image and input image must have the same dimensions.") |
|
|
|
downscaled_image: Optional[ImageRef] = None |
|
quantized_image: Optional[ImageRef] = None |
|
|
|
if output_quantized_path: |
|
quantized_image = ImageRef(Image.new("RGB", input_image.size)) |
|
|
|
extracted_boxes = extract_boxes(control_image) |
|
if output_colorized_full_path: |
|
extracted_boxes.to_colored_checkers(full=True).save(output_colorized_full_path) |
|
if output_colorized_down_path: |
|
extracted_boxes.to_colored_checkers(full=False).save(output_colorized_down_path) |
|
|
|
if output_downscaled_path: |
|
downscaled_image = ImageRef(Image.new("RGB", extracted_boxes.down_dimensions())) |
|
|
|
boxes = extracted_boxes.boxes() |
|
downsample_all(input_image=input_image, output_image=quantized_image, down_image=downscaled_image, boxes=boxes, sample_radius=sample_radius, downsampler=downsampler, trim_cropped_edges=trim_cropped_edges) |
|
|
|
if quantized_image: |
|
assert output_quantized_path |
|
quantized_image.ref.save(output_quantized_path) |
|
if downscaled_image: |
|
assert output_downscaled_path |
|
downscaled_image.ref.save(output_downscaled_path) |
|
|
|
|
|
def main(cli_args: List[str]) -> None: |
|
parser = argparse.ArgumentParser(description="Downsample and rescale image.") |
|
parser.add_argument("--control", type=str, required=True, help="Path to control image.") |
|
parser.add_argument("--input", type=str, required=True, help="Path to input image.") |
|
parser.add_argument("--output-downscaled", type=str, help="Path to save the output downscaled image.") |
|
parser.add_argument("--output-quantized", type=str, help="Path to save the output quantized image (downscaled and then upscaled to the original size).") |
|
parser.add_argument("--sample-radius", type=int, default=None, help="Radius for sampling (Manhattan distance).") |
|
parser.add_argument("--downsampler", choices=["box", "bilinear", "bicubic", "hamming", "lanczos"], default="box", help="Downsampler to use.") |
|
parser.add_argument("--trim-cropped-edges", type=str2bool, default=False, help="Drop mapped checker grid elements that are cropped in the control image.") |
|
parser.add_argument("--output-colorized-full", type=str, help="Colorize the full checker image to debug the checker parsing.") |
|
parser.add_argument("--output-colorized-down", type=str, help="Colorize the downscaled checker image to debug the checker parsing.") |
|
|
|
args = parser.parse_args(cli_args) |
|
downsampler = Image.Resampling[args.downsampler.upper()] |
|
|
|
controlled_downscale( |
|
control_path=args.control, |
|
input_path=args.input, |
|
output_downscaled_path=args.output_downscaled, |
|
output_quantized_path=args.output_quantized, |
|
sample_radius=args.sample_radius, |
|
downsampler=downsampler, |
|
trim_cropped_edges=args.trim_cropped_edges, |
|
output_colorized_full_path=args.output_colorized_full, |
|
output_colorized_down_path=args.output_colorized_down, |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
main(sys.argv[1:]) |
|
|
