gpu

app.py

@@ -8,6 +8,7 @@ import numpy as np
 from PIL import Image
 import base64
 from io import BytesIO
+import os
 
 import dataset
 from dataset import MyDataset, ImageKeypointDataset, load_filenames, load_keypoints
@@ -92,8 +93,7 @@ model_index = 0
 
 # ヒートマップの生成関数
 @spaces.GPU
-def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
-    global model_index, mean_vector_list
+def get_heatmaps(model_info, source_num, x_coords, y_coords, uploaded_image):
     if type(uploaded_image) == str:
         uploaded_image = Image.open(uploaded_image)
     if type(source_num) == str:
@@ -102,6 +102,13 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
         x_coords = int(x_coords)
     if type(y_coords) == str:
         y_coords = int(y_coords)
+        
+    if type(model_info) == str:
+        model_info = eval(model_info)
+    model_index = models_info.index(model_info)
+
+    mean_vector_list = np.load(f"resources/mean_vector_list_{model_info['name']}.npy", allow_pickle=True)
+    mean_vector_list = torch.tensor(mean_vector_list).to(device)
 
     dec5, _ = models[model_index](x)
     feature_map = dec5
@@ -138,24 +145,6 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
     plt.close(fig)
     return fig
 
-@spaces.GPU
-def setup(model_info, input_image=None):
-    global model_index, mean_vector_list
-    # str -> dictに変換
-    if type(model_info) == str:
-        model_info = eval(model_info)
-
-    model_index = models_info.index(model_info)
-
-    feature_map, _ = models[model_index](test_imgs)
-    mean_vector_list = utils.get_mean_vector(feature_map, points)
-    
-    if input_image is not None:
-        fig = get_heatmaps(0, image_size // 2, image_size // 2, input_image)
-        return fig
-    
-    print("setup done.")
-
 with gr.Blocks() as demo:
     # title
     gr.Markdown("# TripletGeoEncoder Feature Map Visualization")
@@ -168,24 +157,17 @@ with gr.Blocks() as demo:
                 "For further information, please contact me on X (formerly Twitter): @Yeq6X.")
     
     gr.Markdown("## Heatmap Visualization")
-
+ 
+    model_info = gr.Dropdown(
+        choices=[str(model_info) for model_info in models_info],
+        container=False
+    )
     input_image = gr.ImageEditor(label="Cropped Image", elem_id="input_image", crop_size=(112, 112), show_fullscreen_button=True)
-    output_plot = gr.Plot(value=None, elem_id="output_plot", show_label=False)
-    with gr.Row():
-        with gr.Column():
-            with gr.Row():
-                model_name = gr.Dropdown(
-                    choices=[str(model_info) for model_info in models_info],
-                    container=False
-                )
-            load_button = gr.Button("Load Model")
-            load_button.click(setup, inputs=[model_name, input_image], outputs=[output_plot])
-        with gr.Row():
-            pass    
-
+    output_plot = gr.Plot(value=None, elem_id="output_plot", show_label=False)   
     inference = gr.Interface(
         get_heatmaps,
         inputs=[
+            model_info,
             gr.Slider(0, batch_size - 1, step=1, label="Source Image Index"),
             gr.Slider(0, image_size - 1, step=1, value=image_size // 2, label="X Coordinate"),
             gr.Slider(0, image_size - 1, step=1, value=image_size // 2, label="Y Coordinate"),
@@ -205,8 +187,5 @@ with gr.Blocks() as demo:
         inputs=[input_image],
     )
 
-    setup(models_info[0])
-    print(mean_vector_list)
-
     demo.launch()
 

resources/mean_vector_list_ae_model_tf_2024-03-05_00-35-21.pth.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a8f81b924edda413139a5743408c6f38ebd7930b5d39cc98a6b4dd49bd42dae
+size 3328

resources/mean_vector_list_autoencoder-epoch=09-train_loss=1.00.ckpt.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21c0d7cd81ee6e7fac9f0333209daec9e74a7c1e72e358b7732e8ecb3efea5f2
+size 6528

resources/mean_vector_list_autoencoder-epoch=29-train_loss=1.01.ckpt.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9edb6ad2b6a0b9121905ea04ac1a39618d709329045c9cf00673f6281fc412c
+size 6528

resources/mean_vector_list_autoencoder-epoch=49-train_loss=1.01.ckpt.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4141deceac789a26c3d6cf02f20068e8caf37fb14b8af3f63ae365b32462d76f
+size 6528

utils.py

@@ -132,7 +132,7 @@ def get_mean_vector(feature_map, points):
       x_coords, y_coords = torch.round(points[:,i].t()).to(torch.long)
       vectors = feature_map[torch.arange(feature_map.size(0)), :, y_coords, x_coords]  # 1次元ベクトルに合わせてサイズを調整
       # mean_vector = vectors[0:10].mean(0) # 10個の特徴マップの平均ベクトルを取得
-      mean_vector = vectors.mean(0)
+      mean_vector = vectors.mean(0).detach().cpu().numpy()
       mean_vector_list.append(mean_vector)
   return mean_vector_list