add description and examples

.gitignore

@@ -0,0 +1,2 @@
+venv
+__pycache__/

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import spaces
+# import spaces
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
@@ -14,17 +14,17 @@ from io import BytesIO
 
 # モデルとデータの読み込み
 def load_model():
-    model_path = "checkpoints/ae_model_tf_2024-03-05_00-35-21.pth"
-    feature_dim = 32
+    model_path = "checkpoints/autoencoder-epoch=49-train_loss=1.01.ckpt"
+    feature_dim = 64
     model = AutoencoderModule(feature_dim=feature_dim)
     state_dict = torch.load(model_path)
     
-    # state_dict のキーを修正
-    new_state_dict = {}
-    for key in state_dict:
-        new_key = "model." + key
-        new_state_dict[new_key] = state_dict[key]
-    model.load_state_dict(new_state_dict)
+    # # state_dict のキーを修正
+    # new_state_dict = {}
+    # for key in state_dict:
+    #     new_key = "model." + key
+    #     new_state_dict[new_key] = state_dict[key]
+    model.load_state_dict(state_dict['state_dict'])
     model.eval()
     
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -58,7 +58,9 @@ x = load_data(device)
 
 # アップロード画像の前処理
 def preprocess_uploaded_image(uploaded_image, image_size):
-    uploaded_image = Image.fromarray(uploaded_image)
+    # ndarrayの場合はPILイメージに変換
+    if type(uploaded_image) == np.ndarray:
+        uploaded_image = Image.fromarray(uploaded_image)
     uploaded_image = uploaded_image.convert("RGB")
     uploaded_image = uploaded_image.resize((image_size, image_size))
     uploaded_image = np.array(uploaded_image).transpose(2, 0, 1) / 255.0
@@ -66,8 +68,17 @@ def preprocess_uploaded_image(uploaded_image, image_size):
     return uploaded_image
 
 # ヒートマップの生成関数
-@spaces.GPU
+# @spaces.GPU
 def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
+    if type(uploaded_image) == str:
+        uploaded_image = Image.open(uploaded_image)
+    if type(source_num) == str:
+        source_num = int(source_num)
+    if type(x_coords) == str:
+        x_coords = int(x_coords)
+    if type(y_coords) == str:
+        y_coords = int(y_coords)
+    
     with torch.no_grad():
         dec5, _ = model(x)
         img = x
@@ -101,7 +112,7 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
         source_map = norm_batch_distance_map[source_num]
         target_map = norm_batch_distance_map[target_num]
 
-        alpha = 0.8
+        alpha = 0.7
         blended_source = (1 - alpha) * img[source_num] + alpha * torch.cat(((norm_batch_distance_map[source_num] / norm_batch_distance_map[source_num].max()).unsqueeze(0), torch.zeros(2, image_size, image_size, device=device)))
         blended_target = (1 - alpha) * img[target_num] + alpha * torch.cat(((norm_batch_distance_map[target_num] / norm_batch_distance_map[target_num].max()).unsqueeze(0), torch.zeros(2, image_size, image_size, device=device)))
         
@@ -154,6 +165,7 @@ async () => {
             console.log(files);
             if (files && files.length > 0) {
                 console.log("File selected");
+                document.querySelector("#input_file_button").style.display = "none";
                 document.querySelector("#crop_view").style.display = "block";
                 document.querySelector("#crop_button").style.display = "block";
                 const url = URL.createObjectURL(files[0]);
@@ -183,6 +195,7 @@ async () => {
 
                 document.getElementById("crop_view").style.display = "none";
                 document.getElementById("crop_button").style.display = "none";
+                document.querySelector("#input_file_button").style.display = "block";
                 
                 cropper.destroy();
             }
@@ -194,6 +207,16 @@ async () => {
 """
 
 with gr.Blocks() as demo:
+    # title
+    gr.Markdown("# TripletGeoEncoder Feature Map Visualization")
+    # description
+    gr.Markdown("This demo visualizes the feature maps of a TripletGeoEncoder trained on the CelebA dataset using self-supervised learning without annotations from only 1000 images. "
+                  "The feature maps are visualized as heatmaps, where the source map shows the distance of each pixel in the source image to the selected pixel, and the target map shows the distance of each pixel in the target image to the selected pixel. "
+
+                "The blended source and target images show the source and target images with the source and target maps overlaid, respectively. "
+
+                "For further information, please contact me on X (formerly Twitter): @Yeq6X.")
+
     with gr.Row():
         with gr.Column():
             source_num = gr.Slider(0, batch_size - 1, step=1, label="Source Image Index")
@@ -202,20 +225,30 @@ with gr.Blocks() as demo:
 
             # GradioのFileコンポーネントでファイル選択ボタンを追加
             gr.HTML('<input type="file" id="input_file" style="display:none;">')
-            input_file_button = gr.Button("画像を選択", elem_id="input_file_button")
+            input_file_button = gr.Button("Upload Image and Crop", elem_id="input_file_button", variant="primary")
+            crop_button = gr.Button("Crop", elem_id="crop_button", variant="primary")
             # 画像を表示するためのHTML画像タグをGradioで表示
             gr.HTML('<img id="crop_view" style="max-width:100%;">')
             # Gradioのボタンコンポーネントを追加し、IDを付与
-            crop_button = gr.Button("クロップ", elem_id="crop_button", variant="primary")
             # クロップされた画像データのテキストボックス（Base64データ）
             cropped_image_data = gr.Textbox(visible=False, elem_id="cropped_image_data")
-            input_image = gr.Image(label="Cropped Image", interactive=False)
+            input_image = gr.Image(label="Cropped Image", elem_id="input_image")
             # cropped_image_dataが更新されたらprocess_imageを呼び出す
             cropped_image_data.change(process_image, inputs=cropped_image_data, outputs=input_image)
 
+            # examples
+            gr.Markdown("# Examples")
+            gr.Examples(
+                examples=[
+                ["0", "50", "50", "resources/examples/2488.jpg"],
+                ["0", "50", "50", "resources/examples/2899.jpg"]
+            ],
+                inputs=[source_num, x_coords, y_coords, input_image],
+            )
         with gr.Column():
             output_plot = gr.Plot()
 
+
         # Gradioインターフェースの代わり
         source_num.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
         x_coords.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
@@ -223,7 +256,7 @@ with gr.Blocks() as demo:
         input_image.change(get_heatmaps, inputs=[source_num, x_coords, y_coords, input_image], outputs=output_plot)
 
         # JavaScriptコードをロード
-        demo.load(None, None, None, js=scripts)
-
+        demo.load(None, None, None, js=scripts)        
+        
     demo.launch()