def paste(self):
if self.copy_mask is None:
self.show_msg('Please select a region to copy first.')
return
copy_to = renormalize.from_url(self.editing_canvas.mask,
target='pt',
size=(256, 256))[0]
area = renormalize.from_url(self.copy_mask,
target='pt',
size=(256, 256))[0]
t, l, b, r = positive_bounding_box(area)
area = area[t:b, l:r]
target_rgb = self.copy_canvas.rgb
source_rgb = renormalize.from_url(self.editing_canvas.image).permute(
1, 2, 0)
rendered_img = paste_clip_at_center(source_rgb, target_rgb,
centered_location(copy_to),
area)[0].permute(2, 0, 1)
self.editing_canvas.mask = ''
self.editing_canvas.image = renormalize.as_url(rendered_img)
self.positive_masks[self.editing_canvas.index] += copy_to
self.real_images_array[
self.editing_canvas.index].src = renormalize.as_url(
F.interpolate(rendered_img.unsqueeze(dim=0),
size=(self.size, self.size)).squeeze())
self.from_editing_canvas()
def save(self):
if self.editname_textbox.value == '':
self.show_msg('Please enter a name to save your file')
return
savedir = os.path.join(self.savedir, self.editname_textbox.value)
# clear the savedir if conflicting
if os.path.exists(savedir):
for x in os.listdir(savedir):
os.remove(os.path.join(savedir, x))
os.makedirs(savedir, exist_ok=True)
for i in range(self.num_canvases):
if self.real_canvas_array[i].negative_mask != '':
torch.save(self.real_canvas_array[i].negative_mask,
os.path.join(savedir, f'{i}_neg.pt'))
if self.positive_masks[i].sum() != 0:
image = renormalize.from_url(
self.real_canvas_array[i].image) / 2 + 0.5
utils.save_image(image, os.path.join(savedir, f'{i}_rgb.png'))
torch.save(self.positive_masks[i].clamp_(0, 1),
os.path.join(savedir, f'{i}_pos.pt'))
with open(os.path.join(savedir, 'edit_type.txt'), 'w') as f:
f.write(f'{self.edit_type}')
self.show_msg('Done saving')
def copy(self, ev):
self.copy_mask = self.copy_canvas.mask
tgt_style = self.copy_canvas.instance_style
index = self.copy_canvas.index
area = renormalize.from_url(self.copy_mask,
target='pt',
size=(256, 256))[0]
t, l, b, r = positive_bounding_box(area)
H, W, focal = self.hwfs[0]
H, W = H.item(), W.item()
with torch.no_grad():
rays_o, rays_d = get_rays(int(H), int(W), focal, self.poses[index])
rays_o, rays_d = rays_o[t:b, l:r], rays_d[t:b, l:r]
rays_o, rays_d = rays_o.contiguous().view(
-1, rays_o.shape[-1]), rays_d.contiguous().view(
-1, rays_d.shape[-1])
batch_rays = torch.stack([rays_o, rays_d], 0)
# render the rays under the editing canvas color style
style = torch.cat([tgt_style[:, :32], self.instance_style[:, 32:]],
dim=1)
style = style.repeat((batch_rays.shape[1], 1))
rgb, disp, acc, extras = render(H,
W,
focal.item(),
style=style,
rays=batch_rays,
**self.test_kwargs)
self.copy_canvas.rgb = rgb.view(b - t, r - l, -1).cpu() * 2 - 1
self.copy_canvas.mask = ''
def get_image_dataset(self):
images = []
poses = []
positive_masks = []
negative_masks = []
for i in range(self.num_canvases):
# TODO: speed the .sum() up by having an edited field
if self.real_canvas_array[
i].negative_mask != '' or self.positive_masks[i].sum(
) != 0:
image = renormalize.from_url(
self.real_canvas_array[i].image) / 2 + 0.5
if self.real_canvas_array[i].negative_mask != '':
negative_mask = renormalize.from_url(
self.real_canvas_array[i].negative_mask) / 2 + 0.5
negative_mask = F.interpolate(
negative_mask.unsqueeze(dim=0),
size=(self.size * 2, self.size * 2)).squeeze()
negative_masks.append(
(negative_mask > 0).float().clamp_(0, 1))
else:
negative_masks.append(
torch.zeros(self.positive_masks[i].shape).cpu())
if self.positive_masks[i].sum() != 0:
positive_masks.append(self.positive_masks[i].clamp_(0, 1))
else:
positive_masks.append(
torch.zeros(negative_masks[-1].shape).cpu())
images.append(image)
poses.append(self.poses[i])
images = torch.stack(images).permute(0, 2, 3, 1)
positive_masks = torch.stack(positive_masks).permute(0, 2, 3, 1)
negative_masks = torch.stack(negative_masks).permute(0, 2, 3, 1)
poses = torch.stack(poses)
return images, positive_masks, negative_masks, poses
def change_mask(self, ev):
if self.mask_type == 'positive' or self.mask_type == 'sigma':
i = self.editing_canvas.index
orig_img = renormalize.from_url(self.editing_canvas.image)
mask = renormalize.from_url(self.editing_canvas.mask) / 2 + 0.5
mask = F.interpolate(mask.unsqueeze(dim=0),
size=(self.size * 2,
self.size * 2)).squeeze()
if self.mask_type == 'positive':
self.edit_type = 'color'
if self.color is None:
self.show_msg('Please select a color.')
if ev.target.image != '':
self.real_canvas_array[
ev.target.index].negative_mask = ''
return
edited_img = orig_img * (1 - mask) + mask * self.color
elif self.mask_type == 'sigma':
self.edit_type = 'removal'
edited_img = orig_img * (1 - mask) + mask * torch.ones(
(3, self.size * 2, self.size * 2)).to(mask.device)
self.positive_masks[i] += mask
self.real_canvas_array[i].image = renormalize.as_url(edited_img)
self.editing_canvas.image = renormalize.as_url(edited_img)
self.real_images_array[i].src = renormalize.as_url(
F.interpolate(edited_img.unsqueeze(dim=0),
size=(self.size, self.size)).squeeze())
self.editing_canvas.mask = ''
elif self.mask_type == 'negative':
i = ev.target.index
self.real_canvas_array[i].negative_mask = self.editing_canvas.mask
elif self.copy_mask is not None:
self.paste()
else:
if ev.target.image != '':
self.real_canvas_array[ev.target.index].negative_mask = ''
def set_color(self, evt):
for i in range(len(self.color_pallete)):
self.color_pallete[i].style = {
'display': 'inline',
'border': "2px solid white"
}
evt.target.style = {'display': 'inline', 'border': "1px solid black"}
if evt.target.color_type == 'bg':
self.negative_mask()
else:
image = renormalize.from_url(evt.target.src) / 2 + 0.5
image = image * 255
self.color = [int(x) * 2 / 255. - 1 for x in image[:, 0, 0]]
color = torch.zeros((3, self.size * 2, self.size * 2)).cpu()
color[0, :, :] = self.color[0]
color[1, :, :] = self.color[1]
color[2, :, :] = self.color[2]
self.color = color