nan
0cf04a13c6
- Added upscale operator to handle image upscaling by 4x using AI models. - Introduced new jobs for upscale processing in the server, including debug and standard upscale jobs. - Created a new ONNX model handler for upscale operations, supporting tiled inference and resource management. - Updated tests to cover new upscale functionality, including input validation and output verification. - Added documentation for the upscale feature, detailing usage and model options.
590 lines
21 KiB
Python
590 lines
21 KiB
Python
import json
|
|
import math
|
|
|
|
import bpy
|
|
|
|
from ...properties import (
|
|
image_edit_refine_selection_property as refine_selection_property,
|
|
)
|
|
from ...runtime import JobOperatorBase
|
|
from ..balloon.operator import (
|
|
deserialize_matrix,
|
|
drawing_in_region,
|
|
hidden_string_property,
|
|
screen_path_to_image_pixels,
|
|
serialize_matrix,
|
|
simplify_screen_path,
|
|
)
|
|
from ..common import HIDDEN, cleanup_temporary_input, push_undo
|
|
from ..image import (
|
|
is_image_empty,
|
|
plane_bounds,
|
|
prepare_image_input,
|
|
remove_temporary_input,
|
|
video_total_frames,
|
|
)
|
|
from .common import (
|
|
active_view3d_tool_id,
|
|
apply_image_mask_settings,
|
|
crop_pixels_to_alpha,
|
|
draw_dashed_line,
|
|
draw_image_mask_operator_settings,
|
|
draw_polygon_fill,
|
|
ImageEditWarning,
|
|
image_edit_poll,
|
|
image_mask_refine_request,
|
|
polygon_bounds,
|
|
replace_empty_image_pixels,
|
|
report_image_edit_exception,
|
|
resolve_image_edit_click,
|
|
)
|
|
|
|
|
|
GESTURE_TOOL_IDS = {
|
|
"BOX": "anyimage.image_box",
|
|
"LASSO": "anyimage.image_lasso",
|
|
"LINE": "anyimage.image_line",
|
|
"POLYLINE": "anyimage.image_polyline",
|
|
}
|
|
LINE_SNAP_RADIANS = math.radians(5.0)
|
|
POLYLINE_CLOSE_RADIUS = 8.0
|
|
|
|
|
|
def points_inside_polygon(points, polygon):
|
|
import numpy as np
|
|
|
|
points = np.asarray(points, dtype=np.float64)
|
|
polygon = np.asarray(polygon, dtype=np.float64)
|
|
inside = np.zeros(len(points), dtype=bool)
|
|
previous = polygon[-1]
|
|
for current in polygon:
|
|
crosses = (current[1] > points[:, 1]) != (previous[1] > points[:, 1])
|
|
edge_x = (
|
|
(previous[0] - current[0])
|
|
* (points[:, 1] - current[1])
|
|
/ (previous[1] - current[1] + 1e-300)
|
|
+ current[0]
|
|
)
|
|
inside ^= crosses & (points[:, 0] < edge_x)
|
|
previous = current
|
|
return inside
|
|
|
|
|
|
def polygon_mask(image_size, polygon):
|
|
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")
|
|
polygon_bounds(polygon)
|
|
y, x = np.indices((height, width), dtype=np.float64)
|
|
centers = np.column_stack((x.ravel() + 0.5, y.ravel() + 0.5))
|
|
return points_inside_polygon(centers, polygon).reshape((height, width)).astype(
|
|
np.float32
|
|
)
|
|
|
|
|
|
def refined_mask_from_pixels(mask_pixels, mask_size, bounds, source_size):
|
|
import numpy as np
|
|
|
|
mask_width, mask_height = (int(value) for value in mask_size)
|
|
source_width, source_height = (int(value) for value in source_size)
|
|
local = np.flipud(
|
|
np.asarray(mask_pixels, dtype=np.float32).reshape(
|
|
(mask_height, mask_width, 4)
|
|
)
|
|
)[:, :, 3]
|
|
result = np.zeros((source_height, source_width), dtype=np.float32)
|
|
left, top, right, bottom = (int(value) for value in bounds)
|
|
source_left = max(left, 0)
|
|
source_top = max(top, 0)
|
|
source_right = min(right, source_width)
|
|
source_bottom = min(bottom, source_height)
|
|
if source_right <= source_left or source_bottom <= source_top:
|
|
return result
|
|
result[source_top:source_bottom, source_left:source_right] = local[
|
|
source_top - top : source_bottom - top,
|
|
source_left - left : source_right - left,
|
|
]
|
|
return result
|
|
|
|
|
|
def edit_image_mask_pixels(source_pixels, source_size, selection_mask, invert=False):
|
|
"""Apply a mask while preserving source RGB, then tightly crop the canvas."""
|
|
import numpy as np
|
|
|
|
width, height = (int(value) for value in source_size)
|
|
rgba = np.flipud(
|
|
np.asarray(source_pixels, dtype=np.float32).reshape((height, width, 4))
|
|
).copy()
|
|
mask = np.asarray(selection_mask, dtype=np.float32)
|
|
if mask.shape != (height, width):
|
|
raise ValueError("The refined selection does not match the source image")
|
|
mask = np.clip(mask, 0.0, 1.0)
|
|
if invert:
|
|
mask = 1.0 - mask
|
|
rgba[:, :, 3] *= mask
|
|
return crop_pixels_to_alpha(np.flipud(rgba).ravel(), (0, 0, width, height))
|
|
|
|
|
|
def edit_image_object(source_object, selection_mask, invert=False):
|
|
import numpy as np
|
|
|
|
image = source_object.data
|
|
source_size = tuple(image.size)
|
|
source_pixels = np.empty(len(image.pixels), dtype=np.float32)
|
|
image.pixels.foreach_get(source_pixels)
|
|
pixels, crop_bounds = edit_image_mask_pixels(
|
|
source_pixels, source_size, selection_mask, invert
|
|
)
|
|
left, top, right, bottom = crop_bounds
|
|
return replace_empty_image_pixels(
|
|
source_object,
|
|
pixels,
|
|
(right - left, bottom - top),
|
|
crop_bounds,
|
|
)
|
|
|
|
|
|
def gesture_property():
|
|
return bpy.props.EnumProperty(
|
|
items=(
|
|
("BOX", "Box", "Edit with a rectangular gesture"),
|
|
("LASSO", "Lasso", "Edit with a freehand gesture"),
|
|
("LINE", "Line", "Edit one side of a line"),
|
|
("POLYLINE", "Polyline", "Edit with a clicked polygon"),
|
|
),
|
|
default="BOX",
|
|
options={"HIDDEN", "SKIP_SAVE"},
|
|
)
|
|
|
|
|
|
def box_screen_path(start, end):
|
|
x0, y0 = start
|
|
x1, y1 = end
|
|
if abs(x1 - x0) < 2.0 or abs(y1 - y0) < 2.0:
|
|
return ()
|
|
return ((x0, y0), (x1, y0), (x1, y1), (x0, y1))
|
|
|
|
|
|
def snapped_line_endpoint(start, end, snap):
|
|
if not snap:
|
|
return tuple(float(value) for value in end)
|
|
delta_x = float(end[0]) - float(start[0])
|
|
delta_y = float(end[1]) - float(start[1])
|
|
length = math.hypot(delta_x, delta_y)
|
|
if length <= 0.0:
|
|
return tuple(float(value) for value in end)
|
|
angle = round(math.atan2(delta_y, delta_x) / LINE_SNAP_RADIANS)
|
|
angle *= LINE_SNAP_RADIANS
|
|
return (
|
|
float(start[0]) + math.cos(angle) * length,
|
|
float(start[1]) + math.sin(angle) * length,
|
|
)
|
|
|
|
|
|
def line_half_plane_polygon(start, end, rectangle_size):
|
|
"""Clip a rectangle to the left side of one directed line."""
|
|
width, height = (float(value) for value in rectangle_size)
|
|
if math.dist(start, end) < 2.0:
|
|
return ()
|
|
vertices = [(0.0, 0.0), (width, 0.0), (width, height), (0.0, height)]
|
|
|
|
def side(point):
|
|
return (end[0] - start[0]) * (point[1] - start[1]) - (
|
|
end[1] - start[1]
|
|
) * (point[0] - start[0])
|
|
|
|
result = []
|
|
previous = vertices[-1]
|
|
previous_side = side(previous)
|
|
for current in vertices:
|
|
current_side = side(current)
|
|
if (current_side >= 0.0) != (previous_side >= 0.0):
|
|
factor = previous_side / (previous_side - current_side)
|
|
result.append(
|
|
(
|
|
previous[0] + (current[0] - previous[0]) * factor,
|
|
previous[1] + (current[1] - previous[1]) * factor,
|
|
)
|
|
)
|
|
if current_side >= 0.0:
|
|
result.append(current)
|
|
previous = current
|
|
previous_side = current_side
|
|
return tuple(result)
|
|
|
|
|
|
def image_line_half_plane_polygon(line, image_size):
|
|
"""Keep the image-space half-plane shown by the screen-space preview."""
|
|
if len(line) != 2:
|
|
raise ValueError("Drag a longer line across the image plane")
|
|
return line_half_plane_polygon(line[1], line[0], image_size)
|
|
|
|
|
|
class EditImageMask(bpy.types.Operator):
|
|
bl_idname = "anyimage.edit_image_mask"
|
|
bl_label = "Edit Image Mask"
|
|
bl_description = "Edit and tightly crop an Image Empty with a gesture"
|
|
bl_options = {"REGISTER", "UNDO"}
|
|
|
|
source_object_name: hidden_string_property()
|
|
source_matrix_data: hidden_string_property()
|
|
selection_data: hidden_string_property()
|
|
refine_selection: refine_selection_property()
|
|
invert: bpy.props.BoolProperty(
|
|
name="Invert",
|
|
description="Keep the opposite side of the refined gesture mask",
|
|
default=False,
|
|
options={"SKIP_SAVE"},
|
|
)
|
|
gesture: gesture_property()
|
|
|
|
_path = None
|
|
_cursor = None
|
|
_handle = None
|
|
_area_pointer = 0
|
|
_region_pointer = 0
|
|
_region_size = (0, 0)
|
|
|
|
def draw(self, _context):
|
|
draw_image_mask_operator_settings(self.layout, self)
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
return image_edit_poll(context)
|
|
|
|
def _screen_polygon(self):
|
|
if self.gesture == "BOX":
|
|
return box_screen_path(self._path[0], self._cursor)
|
|
if self.gesture == "LASSO":
|
|
return tuple(simplify_screen_path(self._path))
|
|
if self.gesture == "LINE":
|
|
return line_half_plane_polygon(
|
|
self._path[0], self._cursor, self._region_size
|
|
)
|
|
vertices = list(self._path)
|
|
if self._cursor is not None:
|
|
vertices.append(self._cursor)
|
|
return tuple(vertices)
|
|
|
|
def _draw_selection(self):
|
|
if not self._path or not drawing_in_region(
|
|
self._area_pointer, self._region_pointer
|
|
):
|
|
return
|
|
polygon = self._screen_polygon()
|
|
if len(polygon) >= 3:
|
|
draw_polygon_fill(
|
|
polygon,
|
|
(1.0, 0.25, 0.2, 0.16)
|
|
if self.invert
|
|
else (1.0, 1.0, 1.0, 0.16),
|
|
)
|
|
if self.gesture == "LINE":
|
|
draw_dashed_line((self._path[0], self._cursor))
|
|
return
|
|
if len(polygon) > 1:
|
|
draw_dashed_line((*polygon, polygon[0]))
|
|
if (
|
|
self.gesture == "POLYLINE"
|
|
and len(self._path) >= 3
|
|
and self._cursor is not None
|
|
and math.dist(self._cursor, self._path[0]) <= POLYLINE_CLOSE_RADIUS
|
|
):
|
|
self._draw_polyline_close_indicator()
|
|
|
|
def _draw_polyline_close_indicator(self):
|
|
import gpu
|
|
from gpu_extras.batch import batch_for_shader
|
|
|
|
center_x, center_y = self._path[0]
|
|
vertices = tuple(
|
|
(
|
|
center_x + math.cos(index * math.tau / 24.0) * POLYLINE_CLOSE_RADIUS,
|
|
center_y + math.sin(index * math.tau / 24.0) * POLYLINE_CLOSE_RADIUS,
|
|
)
|
|
for index in range(25)
|
|
)
|
|
shader = gpu.shader.from_builtin("UNIFORM_COLOR")
|
|
batch = batch_for_shader(shader, "LINE_STRIP", {"pos": vertices})
|
|
gpu.state.blend_set("ALPHA")
|
|
gpu.state.line_width_set(2.0)
|
|
shader.bind()
|
|
shader.uniform_float("color", (1.0, 1.0, 1.0, 0.95))
|
|
batch.draw(shader)
|
|
gpu.state.line_width_set(1.0)
|
|
gpu.state.blend_set("NONE")
|
|
|
|
def _set_status(self, context):
|
|
invert = "On" if self.invert else "Off"
|
|
if self.gesture == "POLYLINE":
|
|
action = "click points; click start or Enter to finish; Backspace removes"
|
|
elif self.gesture == "LINE":
|
|
action = "drag line; Ctrl snaps 5°; release to edit"
|
|
else:
|
|
action = "drag and release to edit"
|
|
context.workspace.status_text_set(
|
|
f"Image {self.gesture.title()}: {action}; F toggles Invert ({invert}); "
|
|
"RMB or Esc cancels"
|
|
)
|
|
|
|
def _finish(self, context):
|
|
if self._handle is not None:
|
|
bpy.types.SpaceView3D.draw_handler_remove(self._handle, "WINDOW")
|
|
self._handle = None
|
|
context.workspace.status_text_set(None)
|
|
if context.area is not None:
|
|
context.area.tag_redraw()
|
|
|
|
def cancel(self, context):
|
|
self._finish(context)
|
|
|
|
def invoke(self, context, event):
|
|
source_object, should_edit = resolve_image_edit_click(
|
|
context, event, self.report
|
|
)
|
|
if not should_edit:
|
|
return {"FINISHED"} if source_object is not None else {"CANCELLED"}
|
|
if source_object is None:
|
|
return {"CANCELLED"}
|
|
if video_total_frames(source_object.data) > 1:
|
|
self.report({"WARNING"}, "Image tools only support a still image")
|
|
return {"CANCELLED"}
|
|
apply_image_mask_settings(self, context)
|
|
self.invert = False
|
|
self.source_object_name = source_object.name
|
|
self.source_matrix_data = serialize_matrix(source_object.matrix_world)
|
|
self.selection_data = ""
|
|
point = (float(event.mouse_region_x), float(event.mouse_region_y))
|
|
self._path = [point]
|
|
self._cursor = point
|
|
self._area_pointer = context.area.as_pointer()
|
|
self._region_pointer = context.region.as_pointer()
|
|
self._region_size = (context.region.width, context.region.height)
|
|
self._handle = bpy.types.SpaceView3D.draw_handler_add(
|
|
self._draw_selection, (), "WINDOW", "POST_PIXEL"
|
|
)
|
|
self._set_status(context)
|
|
context.window_manager.modal_handler_add(self)
|
|
context.area.tag_redraw()
|
|
return {"RUNNING_MODAL"}
|
|
|
|
def _complete(self, context):
|
|
source_object = bpy.data.objects.get(self.source_object_name)
|
|
if not is_image_empty(source_object):
|
|
raise RuntimeError("The source Image Empty is no longer available")
|
|
matrix = deserialize_matrix(self.source_matrix_data)
|
|
if self.gesture == "LINE":
|
|
line = screen_path_to_image_pixels(
|
|
(self._path[0], self._cursor),
|
|
context.region,
|
|
context.region_data,
|
|
matrix,
|
|
plane_bounds(source_object),
|
|
tuple(source_object.data.size),
|
|
)
|
|
polygon = image_line_half_plane_polygon(
|
|
line, tuple(source_object.data.size)
|
|
)
|
|
else:
|
|
path = self._screen_polygon()
|
|
if len(path) < 3:
|
|
self._finish(context)
|
|
return {"CANCELLED"}
|
|
polygon = screen_path_to_image_pixels(
|
|
path,
|
|
context.region,
|
|
context.region_data,
|
|
matrix,
|
|
plane_bounds(source_object),
|
|
tuple(source_object.data.size),
|
|
)
|
|
polygon_bounds(polygon)
|
|
self.selection_data = json.dumps(polygon, separators=(",", ":"))
|
|
self._finish(context)
|
|
return self.execute(context)
|
|
|
|
def modal(self, context, event):
|
|
active_tool = active_view3d_tool_id(context)
|
|
if active_tool is not None and active_tool != GESTURE_TOOL_IDS[self.gesture]:
|
|
self._finish(context)
|
|
return {"CANCELLED"}
|
|
if event.type in {"ESC", "RIGHTMOUSE"} and event.value == "PRESS":
|
|
self._finish(context)
|
|
return {"CANCELLED"}
|
|
if event.type == "F" and event.value == "PRESS":
|
|
self.invert = not self.invert
|
|
self._set_status(context)
|
|
context.area.tag_redraw()
|
|
return {"RUNNING_MODAL"}
|
|
if self.gesture == "POLYLINE":
|
|
return self._modal_polyline(context, event)
|
|
if event.type == "MOUSEMOVE":
|
|
point = (float(event.mouse_region_x), float(event.mouse_region_y))
|
|
if self.gesture == "LASSO":
|
|
self._path.append(point)
|
|
self._cursor = point
|
|
elif self.gesture == "LINE":
|
|
self._cursor = snapped_line_endpoint(self._path[0], point, event.ctrl)
|
|
else:
|
|
self._cursor = point
|
|
context.area.tag_redraw()
|
|
return {"RUNNING_MODAL"}
|
|
if event.type == "LEFTMOUSE" and event.value == "RELEASE":
|
|
point = (float(event.mouse_region_x), float(event.mouse_region_y))
|
|
if self.gesture == "LASSO":
|
|
self._path.append(point)
|
|
self._cursor = point
|
|
elif self.gesture == "LINE":
|
|
self._cursor = snapped_line_endpoint(self._path[0], point, event.ctrl)
|
|
if math.dist(self._path[0], self._cursor) < 2.0:
|
|
self._finish(context)
|
|
return {"CANCELLED"}
|
|
else:
|
|
self._cursor = point
|
|
try:
|
|
return self._complete(context)
|
|
except (RuntimeError, TypeError, ValueError) as error:
|
|
report_image_edit_exception(self, error)
|
|
self._finish(context)
|
|
return {"CANCELLED"}
|
|
return {"PASS_THROUGH"}
|
|
|
|
def _modal_polyline(self, context, event):
|
|
if event.type == "BACK_SPACE" and event.value == "PRESS":
|
|
if len(self._path) > 1:
|
|
self._path.pop()
|
|
context.area.tag_redraw()
|
|
return {"RUNNING_MODAL"}
|
|
if event.type in {"RET", "NUMPAD_ENTER"} and event.value == "PRESS":
|
|
return self._confirm_polyline(context)
|
|
if event.type == "MOUSEMOVE":
|
|
self._cursor = (
|
|
float(event.mouse_region_x),
|
|
float(event.mouse_region_y),
|
|
)
|
|
context.area.tag_redraw()
|
|
return {"RUNNING_MODAL"}
|
|
if event.type == "LEFTMOUSE" and event.value == "PRESS":
|
|
point = (float(event.mouse_region_x), float(event.mouse_region_y))
|
|
if (
|
|
len(self._path) >= 3
|
|
and math.dist(point, self._path[0]) <= POLYLINE_CLOSE_RADIUS
|
|
):
|
|
return self._confirm_polyline(context)
|
|
if math.dist(point, self._path[-1]) >= 2.0:
|
|
self._path.append(point)
|
|
context.area.tag_redraw()
|
|
return {"RUNNING_MODAL"}
|
|
return {"PASS_THROUGH"}
|
|
|
|
def _confirm_polyline(self, context):
|
|
self._cursor = self._path[0]
|
|
try:
|
|
return self._complete(context)
|
|
except (RuntimeError, TypeError, ValueError) as error:
|
|
report_image_edit_exception(self, error)
|
|
self._finish(context)
|
|
return {"CANCELLED"}
|
|
|
|
def execute(self, context):
|
|
source_object = bpy.data.objects.get(self.source_object_name)
|
|
if not is_image_empty(source_object):
|
|
self.report({"ERROR"}, "The source Image Empty is no longer available")
|
|
return {"CANCELLED"}
|
|
try:
|
|
polygon = json.loads(self.selection_data)
|
|
polygon_bounds(polygon)
|
|
if video_total_frames(source_object.data) > 1:
|
|
raise RuntimeError("Image tools only support a still image")
|
|
if self.refine_selection:
|
|
source_path, delete_input = prepare_image_input(
|
|
source_object.data, context.scene
|
|
)
|
|
result = bpy.ops.anyimage.refine_image_mask(
|
|
input_path=str(source_path),
|
|
delete_input=delete_input,
|
|
source_object_name=self.source_object_name,
|
|
selection_data=self.selection_data,
|
|
invert=self.invert,
|
|
)
|
|
if "CANCELLED" in result:
|
|
if delete_input:
|
|
remove_temporary_input(source_path)
|
|
return {"CANCELLED"}
|
|
return {"FINISHED"}
|
|
edit_image_object(
|
|
source_object,
|
|
polygon_mask(tuple(source_object.data.size), polygon),
|
|
self.invert,
|
|
)
|
|
except (json.JSONDecodeError, RuntimeError, TypeError, ValueError) as error:
|
|
report_image_edit_exception(self, error)
|
|
return {"CANCELLED"}
|
|
return {"FINISHED"}
|
|
|
|
|
|
class RefineImageMask(JobOperatorBase, bpy.types.Operator):
|
|
bl_idname = "anyimage.refine_image_mask"
|
|
bl_label = "Refine Image Mask"
|
|
job_type = "refine-image-mask"
|
|
starting_message = "Refining image selection"
|
|
|
|
input_path: bpy.props.StringProperty(options=HIDDEN)
|
|
delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN)
|
|
source_object_name: hidden_string_property()
|
|
selection_data: hidden_string_property()
|
|
invert: bpy.props.BoolProperty(default=False, options=HIDDEN)
|
|
|
|
def request(self, _context):
|
|
return image_mask_refine_request(self.input_path, self.selection_data)
|
|
|
|
def response(self, _context, result):
|
|
source_object = bpy.data.objects.get(self.source_object_name)
|
|
if not is_image_empty(source_object):
|
|
raise RuntimeError("The source Image Empty is no longer available")
|
|
import numpy as np
|
|
|
|
bounds = tuple(json.loads(result.file("bounds").read_text(encoding="utf-8")))
|
|
refined_image = bpy.data.images.load(
|
|
str(result.file("image")), check_existing=False
|
|
)
|
|
try:
|
|
refined_pixels = np.empty(len(refined_image.pixels), dtype=np.float32)
|
|
refined_image.pixels.foreach_get(refined_pixels)
|
|
mask = refined_mask_from_pixels(
|
|
refined_pixels,
|
|
tuple(refined_image.size),
|
|
bounds,
|
|
tuple(source_object.data.size),
|
|
)
|
|
finally:
|
|
bpy.data.images.remove(refined_image, do_unlink=True)
|
|
try:
|
|
edit_image_object(source_object, mask, self.invert)
|
|
except ImageEditWarning as error:
|
|
self.report({"WARNING"}, str(error))
|
|
return str(error)
|
|
push_undo("Edit Image Mask")
|
|
return "Image selection refined"
|
|
|
|
def cleanup(self):
|
|
cleanup_temporary_input(self.input_path, self.delete_input)
|
|
|
|
|
|
class ActivateImageEditTool(bpy.types.Operator):
|
|
bl_idname = "anyimage.activate_image_edit_tool"
|
|
bl_label = "Image Edit Tool"
|
|
bl_description = "Activate the Image Box tool in the 3D Viewport"
|
|
bl_options = {"INTERNAL"}
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
return image_edit_poll(context)
|
|
|
|
def execute(self, _context):
|
|
return bpy.ops.wm.tool_set_by_id(
|
|
name="anyimage.image_box", space_type="VIEW_3D"
|
|
)
|