Update

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import os
 from collections import defaultdict
 from skimage.color import deltaE_ciede2000, rgb2lab
+import zipfile
 
 def DoG_filter(image, kernel_size=0, sigma=1.0, k_sigma=2.0, gamma=1.5):
     g1 = cv2.GaussianBlur(image, (kernel_size, kernel_size), sigma)
@@ -23,7 +24,6 @@ def binarize_image(image):
     _, binarized = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
     return binarized
 
-
 def process_XDoG(image_path):
     kernel_size=0
     sigma=1.4
@@ -38,12 +38,10 @@ def process_XDoG(image_path):
     final_image = Image.fromarray(binarized_image)
     return final_image
 
-
-
 def replace_color(image, color_1, blur_radius=2):
     data = np.array(image)
     original_shape = data.shape
-    channels = original_shape[2] if len(original_shape) > 2 else 1  # チャンネル数を確認
+    channels = original_shape[2] if len(original_shape) > 2 else 1
     data = data.reshape(-1, channels)
     color_1 = np.array(color_1)
     matches = np.all(data[:, :3] == color_1, axis=1)
@@ -53,7 +51,7 @@ def replace_color(image, color_1, blur_radius=2):
     while np.any(matches):
         new_matches = np.zeros_like(matches)
         match_num = np.sum(matches)
-        for i in range(len(data)):  # Removed tqdm
+        for i in range(len(data)):
             if matches[i]:
                 x, y = divmod(i, original_shape[1])
                 neighbors = [
@@ -134,13 +132,9 @@ def line_color(image, mask, new_color):
     data[mask, :3] = new_color
     return Image.fromarray(data)
 
-
-
-
-def main(image, lineart):
+def process_image(image, lineart):
     if image.mode != 'RGBA':
         image = image.convert('RGBA')
-
     lineart = lineart.point(lambda x: 0 if x < 200 else 255)
     lineart = ImageOps.invert(lineart)
     kernel = np.ones((3, 3), np.uint8)
@@ -154,47 +148,59 @@ def main(image, lineart):
     replace_color_image = replace_color(filled_image, new_color_1, 2).convert('RGB')
     return replace_color_image
 
-
-# Gradioインターフェース用のメイン関数
-def gradio_interface(image_path):
-
-    # 画像を読み込む
-    image = Image.open(image_path).convert('RGBA')
-
-    #元画像の名前を取得
-    image_name = os.path.basename(image_path)
-
-    # 元のアルファチャンネルを保持
-    alpha = image.getchannel('A') if image.mode == 'RGBA' else None
-
-    # 画像を保存し、RGBAモードで再読み込み
-    image.save(image_path)
-    image = Image.open(image_path).convert('RGBA')
-
-    # 画像処理を実行（アルファチャンネルを無視する処理はここで行わない）
-    rgb_image = image.convert('RGB')
-    lineart = process_XDoG(image_path).convert('L')
-    replace_color_image = main(rgb_image, lineart).convert('RGBA')
-
-    # アルファチャンネルを再適用
-    if alpha:
-        replace_color_image.putalpha(alpha)
-    # image_nameをreplace_color_image_pathに反映
-    replace_color_image_path = image_name + 'replace_color.png'
-
-    # 透過pngで保存
-    replace_color_image.save(replace_color_image_path)
-
-    return replace_color_image_path
-
-
-# Gradioアプリを設定し、起動する
-iface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Image(type='filepath', image_mode="RGBA", label="Original Image"),
-    outputs=gr.Image(type='filepath', image_mode="RGBA", label="Processed Image"), 
-    title="LineArt-Removar",
-    description="画像をアップロードして線画を除去します。"
-)
-
-iface.launch()
+def zip_files(zip_files, zip_path):
+    with zipfile.ZipFile(zip_path, 'w') as zipf:
+        for file_path in zip_files:
+            zipf.write(file_path, arcname=os.path.basename(file_path))
+
+class webui:
+    def __init__(self):
+        self.demo = gr.Blocks()
+
+    def main(self, image_path):
+        image = Image.open(image_path).convert('RGBA')
+        image_name = os.path.basename(image_path)
+        alpha = image.getchannel('A') if image.mode == 'RGBA' else None
+        image.save(image_path)
+        image = Image.open(image_path).convert('RGBA')
+        rgb_image = image.convert('RGB')
+        lineart = process_XDoG(image_path).convert('L')
+        replace_color_image = process_image(rgb_image, lineart).convert('RGBA')
+        if alpha:
+            replace_color_image.putalpha(alpha)
+        replace_color_image_path = f"{image_name}_noline.png"
+        replace_color_image.save(replace_color_image_path)
+        lineart_image = lineart.convert('RGBA')
+        lineart_image.putalpha(255 - lineart_image.getchannel('A'))
+        lineart_image_path = f"{image_name}_lineart.png"
+        lineart_image.save(lineart_image_path)
+
+        zip_files_list = [replace_color_image_path, lineart_image_path]
+        zip_path = f"{image_name}.zip"
+        zip_files(zip_files_list, zip_path)
+        outputs = [replace_color_image, lineart_image]
+        return outputs, zip_path
+
+    def launch(self, share):
+        with self.demo:
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(type='filepath', image_mode="RGBA", label="Original Image")
+                    submit = gr.Button(value="Start")
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Tab("output"):
+                            output_0 = gr.Gallery(format="png")
+                        output_file = gr.File()
+            submit.click(
+                self.main, 
+                inputs=[input_image], 
+                outputs=[output_0, output_file]
+            )
+
+        self.demo.queue()
+        self.demo.launch(share=share)
+
+if __name__ == "__main__":
+    ui = webui()
+    ui.launch(share=True)