-
-
Notifications
You must be signed in to change notification settings - Fork 287
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add node to remove border/frame with constant color (#2794)
* Add node to remove border/frame with constant color * crop border
- Loading branch information
1 parent
60a0bb5
commit 5e404e9
Showing
1 changed file
with
190 additions
and
0 deletions.
There are no files selected for viewing
190 changes: 190 additions & 0 deletions
190
backend/src/packages/chaiNNer_standard/image_dimension/crop/crop_border.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,190 @@ | ||
from __future__ import annotations | ||
|
||
from dataclasses import dataclass | ||
from enum import Enum | ||
|
||
import numpy as np | ||
|
||
from nodes.properties.inputs import EnumInput, ImageInput, NumberInput, SliderInput | ||
from nodes.properties.outputs import ImageOutput | ||
from nodes.utils.utils import Padding, Region, get_h_w_c | ||
|
||
from .. import crop_group | ||
|
||
|
||
class SelectMode(Enum): | ||
ALL_SECTIONS = 1 | ||
CENTER_SECTION = 2 | ||
LARGEST_SECTION = 3 | ||
|
||
|
||
@crop_group.register( | ||
schema_id="chainner:image:remove_border", | ||
name="Crop Border", | ||
description=[ | ||
"Remove the border around the content of an image. The border is assumed to have an approximately constant color. This node can also be used to crop images with an approximately constant background color.", | ||
"The color of the border is automatically determined by the median color of the pixels of a 1px border around the image. The *Tolerance* option can be used to adjust the sensitivity of the color matching. A tolerance of 0% means that a pixel must match the border color exactly, while a tolerance of 5% (default) give a bit of leeway to handle e.g. slight compression artifacts.", | ||
], | ||
icon="MdCrop", | ||
inputs=[ | ||
ImageInput(), | ||
SliderInput( | ||
"Tolerance", | ||
default=5, | ||
minimum=0, | ||
maximum=30, | ||
controls_step=1, | ||
precision=1, | ||
unit="%", | ||
), | ||
EnumInput( | ||
SelectMode, "Select", default=SelectMode.ALL_SECTIONS, label_style="inline" | ||
).with_docs( | ||
"Determines which sections of the image will be selected as the output image.", | ||
"To support removing the border of images with captions (or other information in the border), the *Select* option can be used to select which content section of the image will be returned by the node. The available options are as follows:", | ||
"- All Sections: A single crop containing all sections will be returned.\n" | ||
"- Center Section: The section closest to the center of the image will be returned.\n" | ||
"- Largest Section: The largest section will be returned.", | ||
"So if the image has a border and a caption, *All* will return the inner image + caption, *Center*/*Largest* will return only the inner image (assuming the caption isn't larger than the inner image).", | ||
), | ||
NumberInput("Padding", default=0, minimum=0, maximum=1000, unit="px").with_docs( | ||
"Additional padding around the selected section.", | ||
"This can be used to avoid cutting off parts of the image that are close to the border.", | ||
), | ||
], | ||
outputs=[ | ||
ImageOutput( | ||
image_type=""" | ||
let pad = Input3; | ||
Image { | ||
width: min(max(uint, 1 + 2 * pad), Input0.width), | ||
height: min(max(uint, 1 + 2 * pad), Input0.height), | ||
channels: Input0.channels, | ||
} | ||
""", | ||
assume_normalized=True, | ||
), | ||
], | ||
) | ||
def crop_border_node( | ||
img: np.ndarray, tolerance: float, select: SelectMode, padding: int | ||
) -> np.ndarray: | ||
tolerance /= 100 | ||
|
||
h, w, c = get_h_w_c(img) | ||
|
||
# find the border color of the border | ||
border_color = get_border_color(img) | ||
|
||
# figure out which pixels are likely part of the border | ||
diff: np.ndarray = np.abs(img - border_color) | ||
if c > 1: | ||
# make grayscale | ||
diff = np.mean(diff, axis=-1) | ||
is_content = diff > tolerance | ||
|
||
# get crop region crop bounds | ||
crop = get_crop_region(is_content, select) | ||
crop = crop.add_padding(Padding.all(padding)) | ||
crop = crop.intersect(Region(0, 0, w, h)) | ||
|
||
return crop.read_from(img) | ||
|
||
|
||
def get_crop_region(is_content: np.ndarray, select: SelectMode) -> Region: | ||
assert is_content.ndim == 2 | ||
|
||
# 1. crop horizontally | ||
is_content_horizontal = np.any(is_content, axis=0) | ||
section_w = get_inner_section(is_content_horizontal, select) | ||
is_content = is_content[:, section_w.start : section_w.end] | ||
|
||
# 2. crop vertically | ||
is_content_vertical = np.any(is_content, axis=1) | ||
section_h = get_inner_section(is_content_vertical, select) | ||
is_content = is_content[section_h.start : section_h.end, :] | ||
|
||
crop = Region( | ||
section_w.start, | ||
section_h.start, | ||
section_w.length, | ||
section_h.length, | ||
) | ||
|
||
if select != SelectMode.ALL_SECTIONS: | ||
# 3. crop horizontally again to remove any remaining border | ||
is_content_horizontal = np.any(is_content, axis=0) | ||
section_w = get_inner_section(is_content_horizontal, SelectMode.ALL_SECTIONS) | ||
crop = Region( | ||
crop.x + section_w.start, | ||
crop.y, | ||
section_w.length, | ||
crop.height, | ||
) | ||
|
||
return crop | ||
|
||
|
||
def get_border_color(img: np.ndarray): | ||
""" | ||
Returns the median color in the 1px border of the image. | ||
""" | ||
# Get the 1px border of the image | ||
top = img[0, :, ...] | ||
bottom = img[-1, :, ...] | ||
left = img[:, 0, ...] | ||
right = img[:, -1, ...] | ||
border = np.concatenate((top, bottom, left, right), axis=0) | ||
# Get the median color | ||
return np.median(border, axis=0) | ||
|
||
|
||
@dataclass(frozen=True) | ||
class Section: | ||
start: int | ||
end: int | ||
|
||
@property | ||
def length(self) -> int: | ||
return self.end - self.start | ||
|
||
def union(self, other: Section) -> Section: | ||
return Section(min(self.start, other.start), max(self.end, other.end)) | ||
|
||
def distance_to(self, index: int) -> int: | ||
if index < self.start: | ||
return self.start - index | ||
if index >= self.end: | ||
return index - self.end | ||
return 0 | ||
|
||
|
||
def get_inner_section(is_content: np.ndarray, select: SelectMode) -> Section: | ||
assert is_content.ndim == 1 | ||
size = len(is_content) | ||
|
||
# find all content sections in the image | ||
sections: list[Section] = [] | ||
start = None | ||
for i in range(size): | ||
if not is_content[i]: | ||
if start is not None: | ||
sections.append(Section(start, i)) | ||
start = None | ||
elif start is None: | ||
start = i | ||
if start is not None: | ||
sections.append(Section(start, size)) | ||
start = None | ||
if len(sections) == 0: | ||
return Section(0, size) | ||
|
||
# select the relevant section | ||
if select == SelectMode.ALL_SECTIONS: | ||
return sections[0].union(sections[-1]) | ||
if select == SelectMode.CENTER_SECTION: | ||
distances = [section.distance_to(size // 2) for section in sections] | ||
return sections[np.argmin(distances)] | ||
if select == SelectMode.LARGEST_SECTION: | ||
return max(sections, key=lambda section: section.length) |