Commit e672cfb0 authored by AUTOMATIC's avatar AUTOMATIC

rework of callback for #6094

parent 6062c85d
...@@ -39,12 +39,14 @@ def image_grid(imgs, batch_size=1, rows=None): ...@@ -39,12 +39,14 @@ def image_grid(imgs, batch_size=1, rows=None):
cols = math.ceil(len(imgs) / rows) cols = math.ceil(len(imgs) / rows)
params = script_callbacks.ImageGridLoopParams(imgs, cols, rows)
script_callbacks.image_grid_callback(params)
w, h = imgs[0].size w, h = imgs[0].size
grid = Image.new('RGB', size=(cols * w, rows * h), color='black') grid = Image.new('RGB', size=(params.cols * w, params.rows * h), color='black')
for i, img in enumerate(imgs): for i, img in enumerate(params.imgs):
script_callbacks.image_grid_loop_callback(img) grid.paste(img, box=(i % params.cols * w, i // params.cols * h))
grid.paste(img, box=(i % cols * w, i // cols * h))
return grid return grid
......
...@@ -52,8 +52,10 @@ class UiTrainTabParams: ...@@ -52,8 +52,10 @@ class UiTrainTabParams:
class ImageGridLoopParams: class ImageGridLoopParams:
def __init__(self, img): def __init__(self, imgs, cols, rows):
self.img = img self.imgs = imgs
self.cols = cols
self.rows = rows
ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"]) ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])
...@@ -68,7 +70,7 @@ callback_map = dict( ...@@ -68,7 +70,7 @@ callback_map = dict(
callbacks_cfg_denoiser=[], callbacks_cfg_denoiser=[],
callbacks_before_component=[], callbacks_before_component=[],
callbacks_after_component=[], callbacks_after_component=[],
callbacks_image_grid_loop=[], callbacks_image_grid=[],
) )
...@@ -160,12 +162,14 @@ def after_component_callback(component, **kwargs): ...@@ -160,12 +162,14 @@ def after_component_callback(component, **kwargs):
except Exception: except Exception:
report_exception(c, 'after_component_callback') report_exception(c, 'after_component_callback')
def image_grid_loop_callback(component, **kwargs):
for c in callback_map['callbacks_image_grid_loop']: def image_grid_callback(params: ImageGridLoopParams):
for c in callback_map['callbacks_image_grid']:
try: try:
c.callback(component, **kwargs) c.callback(params)
except Exception: except Exception:
report_exception(c, 'image_grid_loop') report_exception(c, 'image_grid')
def add_callback(callbacks, fun): def add_callback(callbacks, fun):
stack = [x for x in inspect.stack() if x.filename != __file__] stack = [x for x in inspect.stack() if x.filename != __file__]
...@@ -269,9 +273,9 @@ def on_after_component(callback): ...@@ -269,9 +273,9 @@ def on_after_component(callback):
add_callback(callback_map['callbacks_after_component'], callback) add_callback(callback_map['callbacks_after_component'], callback)
def on_image_grid_loop(callback): def on_image_grid(callback):
"""register a function to be called inside the image grid loop. """register a function to be called before making an image grid.
The callback is called with one argument: The callback is called with one argument:
- params: ImageGridLoopParams - parameters to be used inside the image grid loop. - params: ImageGridLoopParams - parameters to be used for grid creation. Can be modified.
""" """
add_callback(callback_map['callbacks_image_grid_loop'], callback) add_callback(callback_map['callbacks_image_grid'], callback)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment