图像修复的可视化demo代码

• 做项目的时候需要用到一个windows窗口可视化来展示我们的工作，我们的工作是一个文本指导的人脸图像修复，所以窗口需要包括输入图像，文本指导输入和修复结果，并且提供在输入图像上画mask的功能，使用tkinter来实现，相关代码如下：

# Authro: Wu
# define of gui
# borrow from https://github.com/zsyzzsoft/co-mod-ganimport tkinter as tk
from PIL import Image, ImageTk, ImageDraw
import numpy as np
import cv2
import torch
import os
def adjust_dynamic_range(data, drange_in, drange_out):if drange_in != drange_out:scale = (np.float32(drange_out[1]) - np.float32(drange_out[0])) / (np.float32(drange_in[1]) - np.float32(drange_in[0]))bias = (np.float32(drange_out[0]) - np.float32(drange_in[0]) * scale)data = data * scale + biasreturn dataclass App(tk.Tk):def __init__(self, model, window_size=256):super().__init__()self.state = -1self.window_size = window_sizeself.canvas = tk.Canvas(self, bg='gray', height=self.window_size, width=self.window_size*2+10)self.canvas.bind("<Button-1>", self.L_press)self.canvas.bind("<ButtonRelease-1>", self.L_release)self.canvas.bind("<B1-Motion>", self.L_move)self.canvas.bind("<Button-3>", self.R_press)self.canvas.bind("<ButtonRelease-3>", self.R_release)self.canvas.bind("<B3-Motion>", self.R_move)self.canvas.bind("<Key>", self.key_down)self.canvas.bind("<KeyRelease>", self.key_up)self.canvas.pack()self.canvas.focus_set()self.canvas_image_left = self.canvas.create_image(0, 0, anchor='nw')self.canvas_image_right = self.canvas.create_image(self.window_size + 10, 0, anchor='nw')# 'generate' buttonself.btn_pen = tk.Button(self, text="生成", command=self.generate)self.btn_pen.pack(side="left", padx="10")# 'continue' buttonself.btn_pen = tk.Button(self, text="继续", command=self.switch)self.btn_pen.pack(side="left", padx="10")# textself.text = Noneself.entry01 = tk.Entry(self, textvariable=self.text, width=50)self.entry01.pack()self.model = modelself.new_image()self.display()def generate(self):real = adjust_dynamic_range(self.input_image, [0, 255], [0, 1])# real: np array (1,3,256,256); # mask: np array (1,1,256,256),masked 0, unmasked 1; # text: stringself.text = self.entry01.get()if len(self.text) < 1:self.text = 'a man'# with torch.no_grad():#     self.output_image = self.model.inpaint_(real, self.mask, self.text)  # save to localsave_path = 'D:\我的文件\dd\demo\mask'if not os.path.exists(save_path):os.makedirs(save_path)mask_to_save = np.repeat(np.expand_dims(np.squeeze((1-self.mask)*255),2),3,2)Image.fromarray(mask_to_save).save(os.path.join(save_path, '7532_mask.png'))input_to_save = np.transpose(np.squeeze(self.input_image*np.repeat(self.mask,3,1)+np.repeat((1-self.mask)*255,3,1)), (1, 2, 0))Image.fromarray(input_to_save).save(os.path.join(save_path, '7532_masked.png'))# predict_to_save = np.transpose(np.squeeze(self.output_image), (1, 2, 0))# Image.fromarray(predict_to_save).save(os.path.join(save_path, 'predict.png'))self.display(2)self.mask = np.ones((1, 1, self.resolution, self.resolution), np.uint8)self.mask_history = [self.mask]def switch(self):self.input_image = self.output_imageself.display(1)def new_image(self):################################# load image from sourceimage_path = tk.filedialog.askopenfilename()# self.input_image = Image.open(image_path).convert('RGB').resize((256,256))self.input_image = Image.open(image_path).convert('RGB')# self.input_image = Image.open('./application/eyeglasses/mmceleba/5995.jpg').convert('RGB').resize((256,256))self.input_image = np.expand_dims(np.transpose(np.asarray(self.input_image, dtype=np.uint8), (2,0,1)),axis=0)#######################self.resolution = self.input_image.shape[-1]self.mask = np.ones((1, 1, self.resolution, self.resolution), np.uint8)self.mask_history = [self.mask]def display(self, state=0):if state != self.state:self.last_state = self.stateself.state = stateif self.state == 0: #  paintingimage = self.input_image * self.mask + (1-self.mask)*255image_for_display = np.transpose(image[0, :3], (1, 2, 0))image_for_display_resized = cv2.resize(image_for_display, (self.window_size, self.window_size))self.tkimage = ImageTk.PhotoImage(image=Image.fromarray(image_for_display_resized))self.canvas.itemconfig(self.canvas_image_left, image=self.tkimage)elif self.state == 1: # switchimage = self.input_imageimage_for_display = np.transpose(image[0, :3], (1, 2, 0))image_for_display_resized = cv2.resize(image_for_display, (self.window_size, self.window_size))self.tkimage = ImageTk.PhotoImage(image=Image.fromarray(image_for_display_resized))self.canvas.itemconfig(self.canvas_image_left, image=self.tkimage)elif self.state == 2: # generateimage =  self.output_image image_for_display = np.transpose(image[0, :3], (1, 2, 0))image_for_display_resized = cv2.resize(image_for_display, (self.window_size, self.window_size))self.tkimage_right = ImageTk.PhotoImage(image=Image.fromarray(image_for_display_resized))self.canvas.itemconfig(self.canvas_image_right, image=self.tkimage_right)image =  self.input_image image_for_display = np.transpose(image[0, :3], (1, 2, 0))image_for_display_resized = cv2.resize(image_for_display, (self.window_size, self.window_size))self.tkimage_left = ImageTk.PhotoImage(image=Image.fromarray(image_for_display_resized))self.canvas.itemconfig(self.canvas_image_left, image=self.tkimage_left)def get_pos(self, event):return (int(event.x * self.resolution / self.window_size), int(event.y * self.resolution / self.window_size))def L_press(self, event):self.last_pos = self.get_pos(event)def L_move(self, event):a = self.last_posb = self.get_pos(event)width = 6img = Image.fromarray(self.mask[0, 0])draw = ImageDraw.Draw(img)draw.line([a, b], fill=0, width=width)draw.ellipse((b[0] - width // 2, b[1] - width // 2, b[0] + width // 2, b[1] + width // 2), fill=0)self.mask = np.array(img)[np.newaxis, np.newaxis, ...]self.display()self.last_pos = bdef L_release(self, event):self.L_move(event)self.mask_history.append(self.mask)def R_press(self, event):self.last_pos = self.get_pos(event)def R_move(self, event):a = self.last_posb = self.get_pos(event)self.mask = self.mask_history[-1].copy()self.mask[0, 0, max(min(a[1], b[1]), 0): max(a[1], b[1]), max(min(a[0], b[0]), 0): max(a[0], b[0])] = 0self.display()def R_release(self, event):self.R_move(event)self.mask_history.append(self.mask)def key_down(self, event):if event.keysym == 'z':if len(self.mask_history) > 1:self.mask_history.pop()self.mask = self.mask_history[-1]self.display()def key_up(self, event):if event.keysym in ['1', '2']:self.display(self.last_state)if __name__ == "__main__":# 这里model不能用None，应当定义一个nn.Module对象并load参数用于修复，这里就不给出来model的定义了。app = App(model=None, window_size=256)app.mainloop()