Add mapped_downscale.py

mapped_downscale.py

@@ -0,0 +1,277 @@
+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
+
+    # inclusive
+    def min(self) -> Tuple[int, int]:
+        return self._min
+    
+    # inclusive
+    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())
+
+    # (left, upper, right, lower)
+    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 average_box_dimensions(boxes: List[DownBox]) -> Dim:
+    assert len(boxes) > 0
+    if len(boxes) == 1:
+        return boxes[0].dimensions()
+    if len(boxes) <= 16:
+        # mean
+        width = 0
+        height = 0
+        for box in boxes:
+            width += box.width()
+            height += box.height()
+        return (width // len(boxes), height // len(boxes))
+    # median
+    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
+    # threshold = 8
+    # if avg[0] > threshold and avg[1] > threshold:
+    #     outlier_dist = 2
+    # threshold = 32
+    # if avg[0] > threshold and avg[1] > threshold:
+    #     outlier_dist = 3
+
+    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 ("yes", "true", "t", "y", "1"):
+        return True
+    elif value.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def main(cli_args: List[str]) -> None:
+    parser = argparse.ArgumentParser(description="Downsample and rescale image.")
+    parser.add_argument("--control", required=True, help="Path to control image.")
+    parser.add_argument("--input", required=True, help="Path to input image.")
+    parser.add_argument("--output-up", help="Path to save the output image, upscaled to the original size.")
+    parser.add_argument("--output-down", help="Path to save the output image, kept at the downsampled 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.")
+
+    args = parser.parse_args(cli_args)
+
+    control_image = Image.open(args.control).convert("1")
+    input_image = Image.open(args.input)
+    if control_image.size != input_image.size:
+        raise ValueError("Control image and input image must have the same dimensions.")
+    downsampler = Image.Resampling[args.downsampler.upper()]
+    output_image: Optional[ImageRef] = None
+    down_image: Optional[ImageRef] = None
+    if not args.output_up and not args.output_down:
+        raise ValueError("At least one of --output-up and --output-down must be specified.")
+    if args.output_up:
+        output_image = ImageRef(Image.new("RGB", input_image.size))
+    extracted_boxes = extract_boxes(control_image)
+    if args.output_down:
+        down_image = ImageRef(Image.new("RGB", extracted_boxes.down_dimensions()))
+
+    boxes = extracted_boxes.boxes()
+
+    print(args.trim_cropped_edges)
+
+    downsample_all(input_image=input_image, output_image=output_image, down_image=down_image, boxes=boxes, sample_radius=args.sample_radius, downsampler=downsampler, trim_cropped_edges=args.trim_cropped_edges)
+    if output_image:
+        output_image.ref.save(args.output_up)
+    if down_image:
+        down_image.ref.save(args.output_down)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])