|
import os |
|
import torch |
|
import numpy as np |
|
import imgui |
|
import dnnlib |
|
from gui_utils import imgui_utils |
|
|
|
|
|
|
|
class DragWidget: |
|
def __init__(self, viz): |
|
self.viz = viz |
|
self.point = [-1, -1] |
|
self.points = [] |
|
self.targets = [] |
|
self.is_point = True |
|
self.last_click = False |
|
self.is_drag = False |
|
self.iteration = 0 |
|
self.mode = 'point' |
|
self.r_mask = 50 |
|
self.show_mask = False |
|
self.mask = torch.ones(256, 256) |
|
self.lambda_mask = 20 |
|
self.feature_idx = 5 |
|
self.r1 = 3 |
|
self.r2 = 12 |
|
self.path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '_screenshots')) |
|
self.defer_frames = 0 |
|
self.disabled_time = 0 |
|
|
|
def action(self, click, down, x, y): |
|
if self.mode == 'point': |
|
self.add_point(click, x, y) |
|
elif down: |
|
self.draw_mask(x, y) |
|
|
|
def add_point(self, click, x, y): |
|
if click: |
|
self.point = [y, x] |
|
elif self.last_click: |
|
if self.is_drag: |
|
self.stop_drag() |
|
if self.is_point: |
|
self.points.append(self.point) |
|
self.is_point = False |
|
else: |
|
self.targets.append(self.point) |
|
self.is_point = True |
|
self.last_click = click |
|
|
|
def init_mask(self, w, h): |
|
self.width, self.height = w, h |
|
self.mask = torch.ones(h, w) |
|
|
|
def draw_mask(self, x, y): |
|
X = torch.linspace(0, self.width, self.width) |
|
Y = torch.linspace(0, self.height, self.height) |
|
yy, xx = torch.meshgrid(Y, X) |
|
circle = (xx - x)**2 + (yy - y)**2 < self.r_mask**2 |
|
if self.mode == 'flexible': |
|
self.mask[circle] = 0 |
|
elif self.mode == 'fixed': |
|
self.mask[circle] = 1 |
|
|
|
def stop_drag(self): |
|
self.is_drag = False |
|
self.iteration = 0 |
|
|
|
def set_points(self, points): |
|
self.points = points |
|
|
|
def reset_point(self): |
|
self.points = [] |
|
self.targets = [] |
|
self.is_point = True |
|
|
|
def load_points(self, suffix): |
|
points = [] |
|
point_path = self.path + f'_{suffix}.txt' |
|
try: |
|
with open(point_path, "r") as f: |
|
for line in f.readlines(): |
|
y, x = line.split() |
|
points.append([int(y), int(x)]) |
|
except: |
|
print(f'Wrong point file path: {point_path}') |
|
return points |
|
|
|
@imgui_utils.scoped_by_object_id |
|
def __call__(self, show=True): |
|
viz = self.viz |
|
reset = False |
|
if show: |
|
with imgui_utils.grayed_out(self.disabled_time != 0): |
|
imgui.text('Drag') |
|
imgui.same_line(viz.label_w) |
|
|
|
if imgui_utils.button('Add point', width=viz.button_w, enabled='image' in viz.result): |
|
self.mode = 'point' |
|
|
|
imgui.same_line() |
|
reset = False |
|
if imgui_utils.button('Reset point', width=viz.button_w, enabled='image' in viz.result): |
|
self.reset_point() |
|
reset = True |
|
|
|
imgui.text(' ') |
|
imgui.same_line(viz.label_w) |
|
if imgui_utils.button('Start', width=viz.button_w, enabled='image' in viz.result): |
|
self.is_drag = True |
|
if len(self.points) > len(self.targets): |
|
self.points = self.points[:len(self.targets)] |
|
|
|
imgui.same_line() |
|
if imgui_utils.button('Stop', width=viz.button_w, enabled='image' in viz.result): |
|
self.stop_drag() |
|
|
|
imgui.text(' ') |
|
imgui.same_line(viz.label_w) |
|
imgui.text(f'Steps: {self.iteration}') |
|
|
|
imgui.text('Mask') |
|
imgui.same_line(viz.label_w) |
|
if imgui_utils.button('Flexible area', width=viz.button_w, enabled='image' in viz.result): |
|
self.mode = 'flexible' |
|
self.show_mask = True |
|
|
|
imgui.same_line() |
|
if imgui_utils.button('Fixed area', width=viz.button_w, enabled='image' in viz.result): |
|
self.mode = 'fixed' |
|
self.show_mask = True |
|
|
|
imgui.text(' ') |
|
imgui.same_line(viz.label_w) |
|
if imgui_utils.button('Reset mask', width=viz.button_w, enabled='image' in viz.result): |
|
self.mask = torch.ones(self.height, self.width) |
|
imgui.same_line() |
|
_clicked, self.show_mask = imgui.checkbox('Show mask', self.show_mask) |
|
|
|
imgui.text(' ') |
|
imgui.same_line(viz.label_w) |
|
with imgui_utils.item_width(viz.font_size * 6): |
|
changed, self.r_mask = imgui.input_int('Radius', self.r_mask) |
|
|
|
imgui.text(' ') |
|
imgui.same_line(viz.label_w) |
|
with imgui_utils.item_width(viz.font_size * 6): |
|
changed, self.lambda_mask = imgui.input_int('Lambda', self.lambda_mask) |
|
|
|
self.disabled_time = max(self.disabled_time - viz.frame_delta, 0) |
|
if self.defer_frames > 0: |
|
self.defer_frames -= 1 |
|
viz.args.is_drag = self.is_drag |
|
if self.is_drag: |
|
self.iteration += 1 |
|
viz.args.iteration = self.iteration |
|
viz.args.points = [point for point in self.points] |
|
viz.args.targets = [point for point in self.targets] |
|
viz.args.mask = self.mask |
|
viz.args.lambda_mask = self.lambda_mask |
|
viz.args.feature_idx = self.feature_idx |
|
viz.args.r1 = self.r1 |
|
viz.args.r2 = self.r2 |
|
viz.args.reset = reset |
|
|
|
|
|
|
|
|
