lnyan's picture
Disable Img2Img to avoid OOM
65ade76
import base64
import json
import io
import numpy as np
from PIL import Image
from pyodide import to_js, create_proxy
import gc
from js import (
console,
document,
devicePixelRatio,
ImageData,
Uint8ClampedArray,
CanvasRenderingContext2D as Context2d,
requestAnimationFrame,
update_overlay,
setup_overlay,
window
)
PAINT_SELECTION = "selection"
IMAGE_SELECTION = "canvas"
BRUSH_SELECTION = "eraser"
NOP_MODE = 0
PAINT_MODE = 1
IMAGE_MODE = 2
BRUSH_MODE = 3
def hold_canvas():
pass
def prepare_canvas(width, height, canvas) -> Context2d:
ctx = canvas.getContext("2d")
canvas.style.width = f"{width}px"
canvas.style.height = f"{height}px"
canvas.width = width
canvas.height = height
ctx.clearRect(0, 0, width, height)
return ctx
# class MultiCanvas:
# def __init__(self,layer,width=800, height=600) -> None:
# pass
def multi_canvas(layer, width=800, height=600):
lst = [
CanvasProxy(document.querySelector(f"#canvas{i}"), width, height)
for i in range(layer)
]
return lst
class CanvasProxy:
def __init__(self, canvas, width=800, height=600) -> None:
self.canvas = canvas
self.ctx = prepare_canvas(width, height, canvas)
self.width = width
self.height = height
def clear_rect(self, x, y, w, h):
self.ctx.clearRect(x, y, w, h)
def clear(self,):
self.clear_rect(0, 0, self.canvas.width, self.canvas.height)
def stroke_rect(self, x, y, w, h):
self.ctx.strokeRect(x, y, w, h)
def fill_rect(self, x, y, w, h):
self.ctx.fillRect(x, y, w, h)
def put_image_data(self, image, x, y):
data = Uint8ClampedArray.new(to_js(image.tobytes()))
height, width, _ = image.shape
image_data = ImageData.new(data, width, height)
self.ctx.putImageData(image_data, x, y)
del image_data
# def draw_image(self,canvas, x, y, w, h):
# self.ctx.drawImage(canvas,x,y,w,h)
def draw_image(self,canvas, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight):
self.ctx.drawImage(canvas, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight)
@property
def stroke_style(self):
return self.ctx.strokeStyle
@stroke_style.setter
def stroke_style(self, value):
self.ctx.strokeStyle = value
@property
def fill_style(self):
return self.ctx.strokeStyle
@fill_style.setter
def fill_style(self, value):
self.ctx.fillStyle = value
# RGBA for masking
class InfCanvas:
def __init__(
self,
width,
height,
selection_size=256,
grid_size=64,
patch_size=4096,
test_mode=False,
) -> None:
assert selection_size < min(height, width)
self.width = width
self.height = height
self.display_width = width
self.display_height = height
self.canvas = multi_canvas(5, width=width, height=height)
setup_overlay(width,height)
# place at center
self.view_pos = [patch_size//2-width//2, patch_size//2-height//2]
self.cursor = [
width // 2 - selection_size // 2,
height // 2 - selection_size // 2,
]
self.data = {}
self.grid_size = grid_size
self.selection_size_w = selection_size
self.selection_size_h = selection_size
self.patch_size = patch_size
# note that for image data, the height comes before width
self.buffer = np.zeros((height, width, 4), dtype=np.uint8)
self.sel_buffer = np.zeros((selection_size, selection_size, 4), dtype=np.uint8)
self.sel_buffer_bak = np.zeros(
(selection_size, selection_size, 4), dtype=np.uint8
)
self.sel_dirty = False
self.buffer_dirty = False
self.mouse_pos = [-1, -1]
self.mouse_state = 0
# self.output = widgets.Output()
self.test_mode = test_mode
self.buffer_updated = False
self.image_move_freq = 1
self.show_brush = False
self.scale=1.0
self.eraser_size=32
def reset_large_buffer(self):
self.canvas[2].canvas.width=self.width
self.canvas[2].canvas.height=self.height
# self.canvas[2].canvas.style.width=f"{self.display_width}px"
# self.canvas[2].canvas.style.height=f"{self.display_height}px"
self.canvas[2].canvas.style.display="block"
self.canvas[2].clear()
def draw_eraser(self, x, y):
self.canvas[-2].clear()
self.canvas[-2].fill_style = "#ffffff"
self.canvas[-2].fill_rect(x-self.eraser_size//2,y-self.eraser_size//2,self.eraser_size,self.eraser_size)
self.canvas[-2].stroke_rect(x-self.eraser_size//2,y-self.eraser_size//2,self.eraser_size,self.eraser_size)
def use_eraser(self,x,y):
if self.sel_dirty:
self.write_selection_to_buffer()
self.draw_buffer()
self.canvas[2].clear()
self.buffer_dirty=True
bx0,by0=int(x)-self.eraser_size//2,int(y)-self.eraser_size//2
bx1,by1=bx0+self.eraser_size,by0+self.eraser_size
bx0,by0=max(0,bx0),max(0,by0)
bx1,by1=min(self.width,bx1),min(self.height,by1)
self.buffer[by0:by1,bx0:bx1,:]*=0
self.draw_buffer()
self.draw_selection_box()
def setup_mouse(self):
self.image_move_cnt = 0
def get_mouse_mode():
mode = document.querySelector("#mode").value
if mode == PAINT_SELECTION:
return PAINT_MODE
elif mode == IMAGE_SELECTION:
return IMAGE_MODE
return BRUSH_MODE
def get_event_pos(event):
canvas = self.canvas[-1].canvas
rect = canvas.getBoundingClientRect()
x = (canvas.width * (event.clientX - rect.left)) / rect.width
y = (canvas.height * (event.clientY - rect.top)) / rect.height
return x, y
def handle_mouse_down(event):
self.mouse_state = get_mouse_mode()
if self.mouse_state==BRUSH_MODE:
x,y=get_event_pos(event)
self.use_eraser(x,y)
def handle_mouse_out(event):
last_state = self.mouse_state
self.mouse_state = NOP_MODE
self.image_move_cnt = 0
if last_state == IMAGE_MODE:
self.update_view_pos(0, 0)
if True:
self.clear_background()
self.draw_buffer()
self.reset_large_buffer()
self.draw_selection_box()
gc.collect()
if self.show_brush:
self.canvas[-2].clear()
self.show_brush = False
def handle_mouse_up(event):
last_state = self.mouse_state
self.mouse_state = NOP_MODE
self.image_move_cnt = 0
if last_state == IMAGE_MODE:
self.update_view_pos(0, 0)
if True:
self.clear_background()
self.draw_buffer()
self.reset_large_buffer()
self.draw_selection_box()
gc.collect()
async def handle_mouse_move(event):
x, y = get_event_pos(event)
x0, y0 = self.mouse_pos
xo = x - x0
yo = y - y0
if self.mouse_state == PAINT_MODE:
self.update_cursor(int(xo), int(yo))
if True:
# self.clear_background()
# console.log(self.buffer_updated)
if self.buffer_updated:
self.draw_buffer()
self.buffer_updated = False
self.draw_selection_box()
elif self.mouse_state == IMAGE_MODE:
self.image_move_cnt += 1
if self.image_move_cnt == self.image_move_freq:
self.draw_buffer()
self.canvas[2].clear()
self.draw_selection_box()
self.update_view_pos(int(xo), int(yo))
self.cached_view_pos=tuple(self.view_pos)
self.canvas[2].canvas.style.display="none"
large_buffer=self.data2array(self.view_pos[0]-self.width//2,self.view_pos[1]-self.height//2,min(self.width*2,self.patch_size),min(self.height*2,self.patch_size))
self.canvas[2].canvas.width=large_buffer.shape[1]
self.canvas[2].canvas.height=large_buffer.shape[0]
# self.canvas[2].canvas.style.width=""
# self.canvas[2].canvas.style.height=""
self.canvas[2].put_image_data(large_buffer,0,0)
else:
self.update_view_pos(int(xo), int(yo), False)
self.canvas[1].clear()
self.canvas[1].draw_image(self.canvas[2].canvas,
self.width//2+(self.view_pos[0]-self.cached_view_pos[0]),self.height//2+(self.view_pos[1]-self.cached_view_pos[1]),
self.width,self.height,
0,0,self.width,self.height
)
self.clear_background()
# self.image_move_cnt = 0
elif self.mouse_state == BRUSH_MODE:
self.use_eraser(x,y)
mode = document.querySelector("#mode").value
if mode == BRUSH_SELECTION:
self.draw_eraser(x,y)
self.show_brush = True
elif self.show_brush:
self.canvas[-2].clear()
self.show_brush = False
self.mouse_pos[0] = x
self.mouse_pos[1] = y
self.canvas[-1].canvas.addEventListener(
"mousedown", create_proxy(handle_mouse_down)
)
self.canvas[-1].canvas.addEventListener(
"mousemove", create_proxy(handle_mouse_move)
)
self.canvas[-1].canvas.addEventListener(
"mouseup", create_proxy(handle_mouse_up)
)
self.canvas[-1].canvas.addEventListener(
"mouseout", create_proxy(handle_mouse_out)
)
async def handle_mouse_wheel(event):
x, y = get_event_pos(event)
self.mouse_pos[0] = x
self.mouse_pos[1] = y
console.log(to_js(self.mouse_pos))
if event.deltaY>10:
window.postMessage(to_js(["click","zoom_out", self.mouse_pos[0], self.mouse_pos[1]]),"*")
elif event.deltaY<-10:
window.postMessage(to_js(["click","zoom_in", self.mouse_pos[0], self.mouse_pos[1]]),"*")
return False
self.canvas[-1].canvas.addEventListener(
"wheel", create_proxy(handle_mouse_wheel), False
)
def clear_background(self):
# fake transparent background
h, w, step = self.height, self.width, self.grid_size
stride = step * 2
x0, y0 = self.view_pos
x0 = (-x0) % stride
y0 = (-y0) % stride
if y0>=step:
val0,val1=stride,step
else:
val0,val1=step,stride
# self.canvas.clear()
self.canvas[0].fill_style = "#ffffff"
self.canvas[0].fill_rect(0, 0, w, h)
self.canvas[0].fill_style = "#aaaaaa"
for y in range(y0-stride, h + step, step):
start = (x0 - val0) if y // step % 2 == 0 else (x0 - val1)
for x in range(start, w + step, stride):
self.canvas[0].fill_rect(x, y, step, step)
self.canvas[0].stroke_rect(0, 0, w, h)
def refine_selection(self):
h,w=self.selection_size_h,self.selection_size_w
h=min(h,self.height)
w=min(w,self.width)
self.selection_size_h=h*8//8
self.selection_size_w=w*8//8
self.update_cursor(1,0)
def update_scale(self, scale, mx=-1, my=-1):
self.sync_to_data()
scaled_width=int(self.display_width*scale)
scaled_height=int(self.display_height*scale)
if max(scaled_height,scaled_width)>=self.patch_size*2-128:
return
if scaled_height<=self.selection_size_h or scaled_width<=self.selection_size_w:
return
if mx>=0 and my>=0:
scaled_mx=mx/self.scale*scale
scaled_my=my/self.scale*scale
self.view_pos[0]+=int(mx-scaled_mx)
self.view_pos[1]+=int(my-scaled_my)
self.scale=scale
for item in self.canvas:
item.canvas.width=scaled_width
item.canvas.height=scaled_height
item.clear()
update_overlay(scaled_width,scaled_height)
self.width=scaled_width
self.height=scaled_height
self.data2buffer()
self.clear_background()
self.draw_buffer()
self.update_cursor(1,0)
self.draw_selection_box()
def update_view_pos(self, xo, yo, update=True):
# if abs(xo) + abs(yo) == 0:
# return
if self.sel_dirty:
self.write_selection_to_buffer()
if self.buffer_dirty:
self.buffer2data()
self.view_pos[0] -= xo
self.view_pos[1] -= yo
if update:
self.data2buffer()
# self.read_selection_from_buffer()
def update_cursor(self, xo, yo):
if abs(xo) + abs(yo) == 0:
return
if self.sel_dirty:
self.write_selection_to_buffer()
self.cursor[0] += xo
self.cursor[1] += yo
self.cursor[0] = max(min(self.width - self.selection_size_w, self.cursor[0]), 0)
self.cursor[1] = max(min(self.height - self.selection_size_h, self.cursor[1]), 0)
# self.read_selection_from_buffer()
def data2buffer(self):
x, y = self.view_pos
h, w = self.height, self.width
if h!=self.buffer.shape[0] or w!=self.buffer.shape[1]:
self.buffer=np.zeros((self.height, self.width, 4), dtype=np.uint8)
# fill four parts
for i in range(4):
pos_src, pos_dst, data = self.select(x, y, i)
xs0, xs1 = pos_src[0]
ys0, ys1 = pos_src[1]
xd0, xd1 = pos_dst[0]
yd0, yd1 = pos_dst[1]
self.buffer[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
def data2array(self, x, y, w, h):
# x, y = self.view_pos
# h, w = self.height, self.width
ret=np.zeros((h, w, 4), dtype=np.uint8)
# fill four parts
for i in range(4):
pos_src, pos_dst, data = self.select(x, y, i, w, h)
xs0, xs1 = pos_src[0]
ys0, ys1 = pos_src[1]
xd0, xd1 = pos_dst[0]
yd0, yd1 = pos_dst[1]
ret[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
return ret
def buffer2data(self):
x, y = self.view_pos
h, w = self.height, self.width
# fill four parts
for i in range(4):
pos_src, pos_dst, data = self.select(x, y, i)
xs0, xs1 = pos_src[0]
ys0, ys1 = pos_src[1]
xd0, xd1 = pos_dst[0]
yd0, yd1 = pos_dst[1]
data[ys0:ys1, xs0:xs1, :] = self.buffer[yd0:yd1, xd0:xd1, :]
self.buffer_dirty = False
def select(self, x, y, idx, width=0, height=0):
if width==0:
w, h = self.width, self.height
else:
w, h = width, height
lst = [(0, 0), (0, h), (w, 0), (w, h)]
if idx == 0:
x0, y0 = x % self.patch_size, y % self.patch_size
x1 = min(x0 + w, self.patch_size)
y1 = min(y0 + h, self.patch_size)
elif idx == 1:
y += h
x0, y0 = x % self.patch_size, y % self.patch_size
x1 = min(x0 + w, self.patch_size)
y1 = max(y0 - h, 0)
elif idx == 2:
x += w
x0, y0 = x % self.patch_size, y % self.patch_size
x1 = max(x0 - w, 0)
y1 = min(y0 + h, self.patch_size)
else:
x += w
y += h
x0, y0 = x % self.patch_size, y % self.patch_size
x1 = max(x0 - w, 0)
y1 = max(y0 - h, 0)
xi, yi = x // self.patch_size, y // self.patch_size
cur = self.data.setdefault(
(xi, yi), np.zeros((self.patch_size, self.patch_size, 4), dtype=np.uint8)
)
x0_img, y0_img = lst[idx]
x1_img = x0_img + x1 - x0
y1_img = y0_img + y1 - y0
sort = lambda a, b: ((a, b) if a < b else (b, a))
return (
(sort(x0, x1), sort(y0, y1)),
(sort(x0_img, x1_img), sort(y0_img, y1_img)),
cur,
)
def draw_buffer(self):
self.canvas[1].clear()
self.canvas[1].put_image_data(self.buffer, 0, 0)
def fill_selection(self, img):
self.sel_buffer = img
self.sel_dirty = True
def draw_selection_box(self):
x0, y0 = self.cursor
w, h = self.selection_size_w, self.selection_size_h
if self.sel_dirty:
self.canvas[2].clear()
self.canvas[2].put_image_data(self.sel_buffer, x0, y0)
self.canvas[-1].clear()
self.canvas[-1].stroke_style = "#0a0a0a"
self.canvas[-1].stroke_rect(x0, y0, w, h)
self.canvas[-1].stroke_style = "#ffffff"
offset=round(self.scale) if self.scale>1.0 else 1
self.canvas[-1].stroke_rect(x0 - offset, y0 - offset, w + offset*2, h + offset*2)
self.canvas[-1].stroke_style = "#000000"
self.canvas[-1].stroke_rect(x0 - offset*2, y0 - offset*2, w + offset*4, h + offset*4)
def write_selection_to_buffer(self):
x0, y0 = self.cursor
x1, y1 = x0 + self.selection_size_w, y0 + self.selection_size_h
self.buffer[y0:y1, x0:x1] = self.sel_buffer
self.sel_dirty = False
self.sel_buffer = np.zeros(
(self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
)
self.buffer_dirty = True
self.buffer_updated = True
# self.canvas[2].clear()
def read_selection_from_buffer(self):
x0, y0 = self.cursor
x1, y1 = x0 + self.selection_size_w, y0 + self.selection_size_h
self.sel_buffer = self.buffer[y0:y1, x0:x1]
self.sel_dirty = False
def base64_to_numpy(self, base64_str):
try:
data = base64.b64decode(str(base64_str))
pil = Image.open(io.BytesIO(data))
arr = np.array(pil)
ret = arr
except:
ret = np.tile(
np.array([255, 0, 0, 255], dtype=np.uint8),
(self.selection_size_h, self.selection_size_w, 1),
)
return ret
def numpy_to_base64(self, arr):
out_pil = Image.fromarray(arr)
out_buffer = io.BytesIO()
out_pil.save(out_buffer, format="PNG")
out_buffer.seek(0)
base64_bytes = base64.b64encode(out_buffer.read())
base64_str = base64_bytes.decode("ascii")
return base64_str
def sync_to_data(self):
if self.sel_dirty:
self.write_selection_to_buffer()
self.canvas[2].clear()
self.draw_buffer()
if self.buffer_dirty:
self.buffer2data()
def sync_to_buffer(self):
if self.sel_dirty:
self.canvas[2].clear()
self.write_selection_to_buffer()
self.draw_buffer()
def resize(self,width,height,scale=None,**kwargs):
self.display_width=width
self.display_height=height
for canvas in self.canvas:
prepare_canvas(width=width,height=height,canvas=canvas.canvas)
setup_overlay(width,height)
if scale is None:
scale=1
self.update_scale(scale)
def save(self):
self.sync_to_data()
state={}
state["width"]=self.display_width
state["height"]=self.display_height
state["selection_width"]=self.selection_size_w
state["selection_height"]=self.selection_size_h
state["view_pos"]=self.view_pos[:]
state["cursor"]=self.cursor[:]
state["scale"]=self.scale
keys=list(self.data.keys())
data={}
for key in keys:
if self.data[key].sum()>0:
data[f"{key[0]},{key[1]}"]=self.numpy_to_base64(self.data[key])
state["data"]=data
return json.dumps(state)
def load(self, state_json):
self.reset()
state=json.loads(state_json)
self.display_width=state["width"]
self.display_height=state["height"]
self.selection_size_w=state["selection_width"]
self.selection_size_h=state["selection_height"]
self.view_pos=state["view_pos"][:]
self.cursor=state["cursor"][:]
self.scale=state["scale"]
self.resize(state["width"],state["height"],scale=state["scale"])
for k,v in state["data"].items():
key=tuple(map(int,k.split(",")))
self.data[key]=self.base64_to_numpy(v)
self.data2buffer()
self.display()
def display(self):
self.clear_background()
self.draw_buffer()
self.draw_selection_box()
def reset(self):
self.data.clear()
self.buffer*=0
self.buffer_dirty=False
self.buffer_updated=False
self.sel_buffer*=0
self.sel_dirty=False
self.view_pos = [0, 0]
self.clear_background()
for i in range(1,len(self.canvas)-1):
self.canvas[i].clear()
def export(self):
self.sync_to_data()
xmin, xmax, ymin, ymax = 0, 0, 0, 0
if len(self.data.keys()) == 0:
return np.zeros(
(self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
)
for xi, yi in self.data.keys():
buf = self.data[(xi, yi)]
if buf.sum() > 0:
xmin = min(xi, xmin)
xmax = max(xi, xmax)
ymin = min(yi, ymin)
ymax = max(yi, ymax)
yn = ymax - ymin + 1
xn = xmax - xmin + 1
image = np.zeros(
(yn * self.patch_size, xn * self.patch_size, 4), dtype=np.uint8
)
for xi, yi in self.data.keys():
buf = self.data[(xi, yi)]
if buf.sum() > 0:
y0 = (yi - ymin) * self.patch_size
x0 = (xi - xmin) * self.patch_size
image[y0 : y0 + self.patch_size, x0 : x0 + self.patch_size] = buf
ylst, xlst = image[:, :, -1].nonzero()
if len(ylst) > 0:
yt, xt = ylst.min(), xlst.min()
yb, xb = ylst.max(), xlst.max()
image = image[yt : yb + 1, xt : xb + 1]
return image
else:
return np.zeros(
(self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
)