Update main.py

main.py

@@ -1,58 +1,76 @@
-import argparse
-
-from PIL import Image
+import gradio as gr
+import tensorflow as tf
 import cv2
 import numpy as np
-from preprocess_image import preprocess_image
-import tensorflow as tf
+from PIL import Image
 import neuralgym as ng
 
+from preprocess_image import preprocess_image
 from inpaint_model import InpaintCAModel
 
-parser = argparse.ArgumentParser()
-parser.add_argument('--image', default='', type=str,
-                    help='The filename of image to be completed.')
-parser.add_argument('--output', default='output.png', type=str,
-                    help='Where to write output.')
-parser.add_argument('--watermark_type', default='istock', type=str,
-                    help='The watermark type')
-parser.add_argument('--checkpoint_dir', default='model/', type=str,
-                    help='The directory of tensorflow checkpoint.')
 
-#checkpoint_dir = 'model/'
+# ===== Inpainting function ===== / ===== 画像修復処理関数 =====
+def inpaint_image(input_image, watermark_type, checkpoint_dir):
+    # Convert image from Gradio (PIL format) / Gradioから受け取る画像をPIL形式で処理
+    image = input_image.convert("RGB")
 
+    # Preprocessing / 前処理
+    input_image = preprocess_image(image, watermark_type)
+    if input_image.shape == (0,):
+        return None
 
-if __name__ == "__main__":
+    # Load configuration file / 設定ファイルの読み込み
     FLAGS = ng.Config('inpaint.yml')
-    # ng.get_gpus(1)
-    args, unknown = parser.parse_known_args()
 
-    model = InpaintCAModel()
-    image = Image.open(args.image)
-    input_image = preprocess_image(image, args.watermark_type)
+    # Reset TensorFlow graph / TensorFlowグラフをリセット
     tf.reset_default_graph()
+    model = InpaintCAModel()
 
+    # GPU configuration / GPU設定
     sess_config = tf.ConfigProto()
     sess_config.gpu_options.allow_growth = True
-    if (input_image.shape != (0,)):
-        with tf.Session(config=sess_config) as sess:
-            input_image = tf.constant(input_image, dtype=tf.float32)
-            output = model.build_server_graph(FLAGS, input_image)
-            output = (output + 1.) * 127.5
-            output = tf.reverse(output, [-1])
-            output = tf.saturate_cast(output, tf.uint8)
-            # load pretrained model
-            vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
-            assign_ops = []
-            for var in vars_list:
-                vname = var.name
-                from_name = vname
-                var_value = tf.contrib.framework.load_variable(
-                    args.checkpoint_dir, from_name)
-                assign_ops.append(tf.assign(var, var_value))
-            sess.run(assign_ops)
-            print('Model loaded.')
-            result = sess.run(output)
-            cv2.imwrite(args.output, cv2.cvtColor(
-                result[0][:, :, ::-1], cv2.COLOR_BGR2RGB))
-            print('image saved to {}'.format(args.output))
+
+    # Start TensorFlow session / TensorFlowセッション開始
+    with tf.Session(config=sess_config) as sess:
+        # Create tensor from image / 画像をテンソルに変換
+        input_image_tensor = tf.constant(input_image, dtype=tf.float32)
+        
+        # Build the model graph / モデルグラフを構築
+        output = model.build_server_graph(FLAGS, input_image_tensor)
+        output = (output + 1.) * 127.5
+        output = tf.reverse(output, [-1])
+        output = tf.saturate_cast(output, tf.uint8)
+
+        # Load model variables / モデル変数を読み込み
+        vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
+        assign_ops = []
+        for var in vars_list:
+            from_name = var.name
+            var_value = tf.contrib.framework.load_variable(checkpoint_dir, from_name)
+            assign_ops.append(tf.assign(var, var_value))
+        sess.run(assign_ops)
+
+        print('Model loaded.')  # モデルの読み込み完了
+        result = sess.run(output)
+        result_img = result[0][:, :, ::-1]  # Convert BGR to RGB / BGRからRGBに変換
+
+    # Convert numpy array to PIL image / numpy配列をPIL画像に変換
+    return Image.fromarray(result_img)
+
+
+# ===== Gradio User Interface ===== / ===== Gradioユーザーインターフェース =====
+iface = gr.Interface(
+    fn=inpaint_image,
+    inputs=[
+        gr.Image(label="Input Image / 入力画像", type="pil"),
+        gr.Radio(["istock", "other"], label="Watermark Type / ウォーターマークタイプ", value="istock"),
+        gr.Textbox(label="Checkpoint Directory / チェックポイントディレクトリ", value="model/")
+    ],
+    outputs=gr.Image(label="Inpainted Image / 修復済み画像"),
+    title="Watermark Inpainting Model / ウォーターマーク除去モデル",
+    description="Upload an image to remove the watermark using a TensorFlow model. / TensorFlowモデルを使用してウォーターマークを除去します。",
+)
+
+# Run the app / アプリを起動
+if __name__ == "__main__":
+    iface.launch()