完整建筑转手绘

import tkinter as tk
from tkinter import filedialog, ttk
from PIL import Image, ImageTk
import requests
import base64
import os
import io
import random
import threading
import time


def encode_pil_to_base64(image):
    with io.BytesIO() as output_bytes:
        image.save(output_bytes, format="PNG")
        bytes_data = output_bytes.getvalue()
    return base64.b64encode(bytes_data).decode("utf-8")


def save_decoded_image(b64_image, folder_path, image_name):
    seq = 0
    output_path = os.path.join(folder_path, f"{image_name}.png")
    while os.path.exists(output_path):
        seq += 1
        output_path = os.path.join(folder_path, f"{image_name}({seq}).png")
    with open(output_path, 'wb') as image_file:
        image_file.write(base64.b64decode(b64_image))
    print(f"Image saved to: {output_path}")
    return output_path  # 返回保存的图片路径


class PhotoSketchGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("建筑效果图转彩色手绘")
        self.root.configure(bg='black')
        self.root.state('zoomed')
        self.folder_path = ''  # 初始化 folder_path
        self.timer_running = False  # 初始化计时器运行状态为False
        self.start_time = None  # 初始化计时器开始时间
        self.buttons_y = None  # 初始化按钮的y坐标
        self.setup_ui()

    def setup_ui(self):
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()
        self.canvas_width = screen_width * 0.4
        self.canvas_height = screen_height * 0.8
        self.left_canvas_x = screen_width * 0.25 - self.canvas_width / 2
        self.right_canvas_x = screen_width * 0.75 - self.canvas_width / 2
        self.setup_buttons_and_canvas()
        self.add_intro_labels()

    def setup_buttons_and_canvas(self):
        self.left_frame = tk.Canvas(self.root, width=self.canvas_width, height=self.canvas_height, bg='black',
                                    highlightthickness=0)
        self.left_frame.place(x=self.left_canvas_x, rely=0.5, anchor='w')
        self.left_frame.create_rectangle(2, 2, self.canvas_width - 2, self.canvas_height - 2, outline='white',
                                         dash=(5, 5))

        self.buttons_y = self.left_frame.winfo_y() + self.canvas_height + 95  # 在画布下方95像素处

        self.upload_btn = tk.Button(self.root, text="上传参考图", command=self.upload_image, bg='darkgrey', fg='white',
                                    font=('微软雅黑', 12))
        self.upload_btn.place(x=self.left_canvas_x, y=self.buttons_y, width=(self.canvas_width / 2 - 5), height=30,
                              anchor='nw')

        self.save_dir_btn = tk.Button(self.root, text="选择保存路径", command=self.select_save_folder, bg='darkgrey',
                                      fg='white', font=('微软雅黑', 12))
        self.save_dir_btn.place(x=self.left_canvas_x + self.canvas_width / 2 + 5, y=self.buttons_y,
                                width=(self.canvas_width / 2 - 5), height=30, anchor='nw')

        self.right_frame = tk.Canvas(self.root, width=self.canvas_width, height=self.canvas_height, bg='black',
                                     highlightthickness=0)
        self.right_frame.place(x=self.right_canvas_x, rely=0.5, anchor='w')
        self.right_frame.create_rectangle(2, 2, self.canvas_width - 2, self.canvas_height - 2, outline='white',
                                          dash=(5, 5))

    def add_intro_labels(self):
        intro_label1 = tk.Label(self.root, text="建筑效果图转彩色手绘", bg='black', fg='white',
                                font=('微软雅黑', 18, 'bold'))
        intro_label1.place(x=self.left_canvas_x, y=10)
        intro_label2 = tk.Label(self.root, text="使用介绍: 请在左侧框内上传一张建筑渲染图", bg='black', fg='white',
                                font=('微软雅黑', 10))
        intro_label2.place(x=self.left_canvas_x, y=45)

    def select_save_folder(self):
        folder_selected = filedialog.askdirectory()
        if folder_selected:
            self.folder_path = folder_selected
            print(f"Selected folder: {self.folder_path}")

    def start_timer(self):
        self.timer_running = True
        self.start_time = time.time()
        self.timer_label = tk.Label(self.root, text="生成时间：0秒", bg='black', fg='white', font=('微软雅黑', 10))
        self.timer_label.place(x=self.right_frame.winfo_x(), y=self.buttons_y, anchor='w')
        self.update_timer()

    def update_timer(self):
        if self.timer_running:
            elapsed_time = round(time.time() - self.start_time)
            self.timer_label.config(text=f"生成时间：{elapsed_time}秒")
            self.root.after(1000, self.update_timer)

    def stop_timer(self):
        self.timer_running = False
        elapsed_time = round(time.time() - self.start_time)
        self.timer_label.config(text=f"生成时间：{elapsed_time}秒")

    def upload_image(self):
        self.upload_btn.config(state='disabled')  # 禁用上传按钮以防止多次上传
        filepath = filedialog.askopenfilename()
        if filepath:
            self.start_timer()  # 开始计时
            self.display_image(filepath, self.left_frame, 'image_label', self.left_canvas_x)
            threading.Thread(target=self.call_api_for_sketch, args=(filepath,)).start()  # 在新线程中调用API

    def display_image(self, image_path, canvas, label_attr, canvas_x):
        img = Image.open(image_path)
        img.thumbnail((canvas.winfo_width(), canvas.winfo_height()))
        img_tk = ImageTk.PhotoImage(img)

        if hasattr(self, label_attr):
            getattr(self, label_attr).configure(image=img_tk)
            getattr(self, label_attr).image = img_tk
        else:
            label = tk.Label(self.root, image=img_tk, bg='black')
            setattr(self, label_attr, label)
            getattr(self, label_attr).image = img_tk
            label.place(x=canvas_x, rely=0.5, anchor='w')

        if label_attr == 'generated_image_label':
            self.stop_timer()

    def call_api_for_sketch(self, image_path):
        # 模拟一些处理时间
        time.sleep(1)
        with Image.open(image_path) as img:
            encoded_image = encode_pil_to_base64(img)
        data = {
            "prompt": "<lora:CWG_archisketch_v1:1>,Building,pre sketch,masterpiece,best quality,featuring markers,(3D:0.7)",
            "negative_prompt": "blurry, lower quality, glossy finish,insufficient contrast",
            "init_images": [encoded_image],
            "steps": 30,
            "width": img.width,
            "height": img.height,
            "seed": random.randint(1, 10000000),
            "alwayson_scripts": {
                "ControlNet": {
                    "args": [
                        {
                            "enabled": "true",
                            "pixel_perfect": "true",
                            "module": "canny",
                            "model": "control_v11p_sd15_canny_fp16 [b18e0966]",
                            "weight": 1,
                            "image": encoded_image
                        },
                        {
                            "enabled": "true",
                            "pixel_perfect": "true",
                            "module": "depth",
                            "model": "control_v11f1p_sd15_depth_fp16 [4b72d323]",
                            "weight": 1,
                            "image": encoded_image
                        }
                    ]
                }
            }
        }

        url = "http://127.0.0.1:7860/sdapi/v1/txt2img"
        response = requests.post(url, json=data)
        if response.status_code == 200:
            response_json = response.json()
            saved_image_path = save_decoded_image(response_json['images'][0], self.folder_path, "sketched_image")
            self.root.after(0, self.display_image, saved_image_path, self.right_frame, 'generated_image_label',
                            self.right_canvas_x)
            self.root.after(0, lambda: self.upload_btn.config(
                state='normal'))  # Re-enable the upload button after processing
        else:
            print("Failed to get response from API, status code:", response.status_code)
            self.root.after(0, lambda: self.upload_btn.config(
                state='normal'))  # Re-enable the upload button if the API call fails


if __name__ == "__main__":
    root = tk.Tk()
    app = PhotoSketchGUI(root)
    root.mainloop()