Add image selection refinement functionality with BEN2 integration
- Implemented viewport.py for handling 3D view projections, gestures, and overlay drawing. - Created refine_image_selection.py to manage the refinement of gesture masks and output image bounds. - Added selection module with refine_selection.py to process and crop images based on alpha channel visibility. - Introduced __init__.py for selection module to expose refine_selection function.
This commit is contained in:
232
src/anyimage/common/selection.py
Normal file
232
src/anyimage/common/selection.py
Normal file
@@ -0,0 +1,232 @@
|
||||
"""Selection values, serialization, rasterization, and Alpha cropping."""
|
||||
|
||||
import json
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
class ImageEditWarning(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SelectionPath:
|
||||
points: tuple[tuple[float, float], ...]
|
||||
invert: bool = False
|
||||
fill_rule: str = "EVEN_ODD"
|
||||
|
||||
def __post_init__(self):
|
||||
points = tuple(
|
||||
(float(point[0]), float(point[1]))
|
||||
for point in self.points
|
||||
)
|
||||
selection_path_bounds(points)
|
||||
if self.fill_rule not in {"EVEN_ODD", "NONZERO"}:
|
||||
raise ValueError(f"Unknown Selection Path fill rule: {self.fill_rule}")
|
||||
object.__setattr__(self, "points", points)
|
||||
object.__setattr__(self, "invert", bool(self.invert))
|
||||
|
||||
def to_json(self):
|
||||
return json.dumps(
|
||||
{
|
||||
"points": self.points,
|
||||
"invert": self.invert,
|
||||
"fill_rule": self.fill_rule,
|
||||
},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, value):
|
||||
data = json.loads(value)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("The Selection Path data is invalid")
|
||||
return cls(
|
||||
points=tuple(data.get("points", ())),
|
||||
invert=bool(data.get("invert", False)),
|
||||
fill_rule=str(data.get("fill_rule", "EVEN_ODD")),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SelectionMask:
|
||||
values: object
|
||||
bounds: tuple[int, int, int, int]
|
||||
|
||||
def __post_init__(self):
|
||||
import numpy as np
|
||||
|
||||
values = np.asarray(self.values, dtype=np.float32)
|
||||
bounds = tuple(int(value) for value in self.bounds)
|
||||
if values.ndim != 2:
|
||||
raise ValueError("The selection mask must be two-dimensional")
|
||||
if len(bounds) != 4:
|
||||
raise ValueError("The selection bounds are invalid")
|
||||
left, top, right, bottom = bounds
|
||||
if right <= left or bottom <= top:
|
||||
raise ValueError("The selection bounds are empty")
|
||||
if values.shape != (bottom - top, right - left):
|
||||
raise ValueError("The selection mask does not match its bounds")
|
||||
if not np.isfinite(values).all():
|
||||
raise ValueError("The selection mask contains invalid values")
|
||||
object.__setattr__(self, "values", np.clip(values, 0.0, 1.0))
|
||||
object.__setattr__(self, "bounds", bounds)
|
||||
|
||||
def full_values(self, image_size):
|
||||
import numpy as np
|
||||
|
||||
width, height = (int(value) for value in image_size)
|
||||
result = np.zeros((height, width), dtype=np.float32)
|
||||
left, top, right, bottom = self.bounds
|
||||
source_left = max(left, 0)
|
||||
source_top = max(top, 0)
|
||||
source_right = min(right, width)
|
||||
source_bottom = min(bottom, height)
|
||||
if source_right <= source_left or source_bottom <= source_top:
|
||||
return result
|
||||
result[source_top:source_bottom, source_left:source_right] = self.values[
|
||||
source_top - top : source_bottom - top,
|
||||
source_left - left : source_right - left,
|
||||
]
|
||||
return result
|
||||
|
||||
|
||||
def selection_path_bounds(points):
|
||||
import numpy as np
|
||||
|
||||
points = np.asarray(points, dtype=np.float64)
|
||||
if points.ndim != 2 or points.shape[0] < 3 or points.shape[1] != 2:
|
||||
raise ValueError("The selection does not define a valid image path")
|
||||
if not np.isfinite(points).all():
|
||||
raise ValueError("The selection contains invalid image coordinates")
|
||||
left = int(math.floor(float(points[:, 0].min())))
|
||||
top = int(math.floor(float(points[:, 1].min())))
|
||||
right = int(math.ceil(float(points[:, 0].max())))
|
||||
bottom = int(math.ceil(float(points[:, 1].max())))
|
||||
if right <= left or bottom <= top:
|
||||
raise ValueError("The selection does not define a valid image area")
|
||||
return left, top, right, bottom
|
||||
|
||||
|
||||
def _rasterize_path(image_size, points, fill_rule):
|
||||
import numpy as np
|
||||
|
||||
width, height = (int(value) for value in image_size)
|
||||
if width <= 0 or height <= 0:
|
||||
raise ValueError("The source image has no readable dimensions")
|
||||
left, top, right, bottom = selection_path_bounds(points)
|
||||
left = max(left, 0)
|
||||
top = max(top, 0)
|
||||
right = min(right, width)
|
||||
bottom = min(bottom, height)
|
||||
mask = np.zeros((height, width), dtype=bool)
|
||||
if right <= left or bottom <= top:
|
||||
return mask
|
||||
|
||||
points = np.asarray(points, dtype=np.float64)
|
||||
x = np.arange(left, right, dtype=np.float64) + 0.5
|
||||
for row_start in range(top, bottom, 128):
|
||||
row_end = min(row_start + 128, bottom)
|
||||
y = np.arange(row_start, row_end, dtype=np.float64) + 0.5
|
||||
state = np.zeros(
|
||||
(len(y), len(x)),
|
||||
dtype=bool if fill_rule == "EVEN_ODD" else np.int16,
|
||||
)
|
||||
previous = points[-1]
|
||||
for current in points:
|
||||
crosses = (current[1] > y) != (previous[1] > y)
|
||||
edge_x = (previous[0] - current[0]) * (y - current[1]) / (
|
||||
previous[1] - current[1] + 1e-300
|
||||
) + current[0]
|
||||
selected = crosses[:, None] & (x[None, :] < edge_x[:, None])
|
||||
if fill_rule == "EVEN_ODD":
|
||||
state ^= selected
|
||||
else:
|
||||
winding = 1 if current[1] > previous[1] else -1
|
||||
state += selected * winding
|
||||
previous = current
|
||||
mask[row_start:row_end, left:right] = state != 0
|
||||
return mask
|
||||
|
||||
|
||||
def _antialias_mask(mask):
|
||||
import numpy as np
|
||||
|
||||
values = np.asarray(mask, dtype=np.float32)
|
||||
kernel = np.asarray((0.029, 0.235, 0.472, 0.235, 0.029), dtype=np.float32)
|
||||
radius = len(kernel) // 2
|
||||
horizontal = np.zeros_like(values)
|
||||
padded = np.pad(values, ((0, 0), (radius, radius)), mode="edge")
|
||||
for offset, weight in enumerate(kernel):
|
||||
horizontal += padded[:, offset : offset + values.shape[1]] * weight
|
||||
result = np.zeros_like(values)
|
||||
padded = np.pad(horizontal, ((radius, radius), (0, 0)), mode="edge")
|
||||
for offset, weight in enumerate(kernel):
|
||||
result += padded[offset : offset + values.shape[0]] * weight
|
||||
return result
|
||||
|
||||
|
||||
def rasterize_selection_path(
|
||||
image_size,
|
||||
selection_path,
|
||||
*,
|
||||
source_alpha=None,
|
||||
alpha_threshold=0.0,
|
||||
antialias=False,
|
||||
):
|
||||
import numpy as np
|
||||
|
||||
width, height = (int(value) for value in image_size)
|
||||
mask = _rasterize_path(
|
||||
(width, height),
|
||||
selection_path.points,
|
||||
selection_path.fill_rule,
|
||||
)
|
||||
if antialias:
|
||||
mask = _antialias_mask(mask)
|
||||
else:
|
||||
mask = mask.astype(np.float32)
|
||||
if selection_path.invert:
|
||||
mask = 1.0 - mask
|
||||
if source_alpha is not None:
|
||||
alpha = np.asarray(source_alpha)
|
||||
if alpha.shape != (height, width):
|
||||
raise ValueError("The source Alpha does not match the image selection")
|
||||
mask = np.where(alpha > float(alpha_threshold), mask, 0.0)
|
||||
visible_y, visible_x = np.nonzero(mask > float(alpha_threshold))
|
||||
if not len(visible_x):
|
||||
raise ImageEditWarning("The selection contains no visible image pixels")
|
||||
bounds = (
|
||||
int(visible_x.min()),
|
||||
int(visible_y.min()),
|
||||
int(visible_x.max()) + 1,
|
||||
int(visible_y.max()) + 1,
|
||||
)
|
||||
left, top, right, bottom = bounds
|
||||
return SelectionMask(mask[top:bottom, left:right], bounds)
|
||||
|
||||
|
||||
def crop_pixels_to_alpha(pixels, bounds):
|
||||
"""Crop Blender-order RGBA pixels to their nonzero Alpha bounds."""
|
||||
import numpy as np
|
||||
|
||||
left, top, right, bottom = bounds
|
||||
width = right - left
|
||||
height = bottom - top
|
||||
rgba = np.flipud(
|
||||
np.asarray(pixels, dtype=np.float32).reshape((height, width, 4))
|
||||
)
|
||||
visible_y, visible_x = np.nonzero(rgba[:, :, 3] > 0.0)
|
||||
if not len(visible_x):
|
||||
raise ImageEditWarning("The edit contains no visible image pixels")
|
||||
content_left = int(visible_x.min())
|
||||
content_top = int(visible_y.min())
|
||||
content_right = int(visible_x.max()) + 1
|
||||
content_bottom = int(visible_y.max()) + 1
|
||||
output = rgba[content_top:content_bottom, content_left:content_right].copy()
|
||||
return np.flipud(output).ravel(), (
|
||||
left + content_left,
|
||||
top + content_top,
|
||||
left + content_right,
|
||||
top + content_bottom,
|
||||
)
|
||||
Reference in New Issue
Block a user