Spaces:
Runtime error
Runtime error
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) | |
def stroke_style(self): | |
return self.ctx.strokeStyle | |
def stroke_style(self, value): | |
self.ctx.strokeStyle = value | |
def fill_style(self): | |
return self.ctx.strokeStyle | |
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=h//8*8 | |
w=w//8*8 | |
h=min(h,self.height) | |
w=min(w,self.width) | |
self.selection_size_h=h | |
self.selection_size_w=w | |
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 | |
) | |