gpu

app.py

@@ -19,7 +19,7 @@ image_size = 112
 batch_size = 32
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-def load_model(model_path="checkpoints/autoencoder-epoch=49-train_loss=1.01.ckpt", feature_dim=64):
+def load_model(model_path, feature_dim):
     model = AutoencoderModule(feature_dim=feature_dim)
     state_dict = torch.load(model_path)
     
@@ -39,7 +39,7 @@ def load_model(model_path="checkpoints/autoencoder-epoch=49-train_loss=1.01.ckpt
     print("Model loaded successfully.")
     return model
 
-def load_data(device, img_dir="resources/trainB/", image_size=112, batch_size=256):
+def load_data(img_dir="resources/trainB/", image_size=112, batch_size=256):
     filenames = load_filenames(img_dir)
     train_X = filenames[:1000]
 
@@ -55,7 +55,7 @@ def load_data(device, img_dir="resources/trainB/", image_size=112, batch_size=25
     print("Data loaded successfully.")
     return x
 
-def load_keypoints(device, img_dir="resources/trainB/", image_size=112, batch_size=32):
+def load_keypoints(img_dir="resources/trainB/", image_size=112, batch_size=32):
     filenames = load_filenames(img_dir)
     train_X = filenames[:1000]
     keypoints = dataset.load_keypoints('resources/DataList.json')
@@ -119,31 +119,39 @@ def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
     plt.close(fig)
     return fig
 
-def setup(model_dict, input_image=None):
+def setup(model_info, input_image=None):
     global model, x, test_imgs, points, mean_vector_list
     # str -> dictに変換
-    if type(model_dict) == str:
-        model_dict = eval(model_dict)
-    model_name = model_dict["name"]
-    feature_dim = model_dict["feature_dim"]
-    model_path = f"checkpoints/{model_name}"
-    model = load_model(model_path, feature_dim)
-    x = load_data(device)
-    test_imgs, points = load_keypoints(device)
+    if type(model_info) == str:
+        model_info = eval(model_info)
+        
+    index = models_info.index(model_info)
+    model = models[index]
+    
+    x = load_data()
+    test_imgs, points = load_keypoints()
+    
     feature_map, _ = model(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
-        
 
-models = [{"name": "ae_model_tf_2024-03-05_00-35-21.pth", "feature_dim": 32},
+
+models_info = [{"name": "ae_model_tf_2024-03-05_00-35-21.pth", "feature_dim": 32},
         {"name": "autoencoder-epoch=09-train_loss=1.00.ckpt", "feature_dim": 64},
         {"name": "autoencoder-epoch=29-train_loss=1.01.ckpt", "feature_dim": 64},
         {"name": "autoencoder-epoch=49-train_loss=1.01.ckpt", "feature_dim": 64}]
-
-setup(models[0])
+models = []
+for model_info in models_info:
+    model_name = model_info["name"]
+    feature_dim = model_info["feature_dim"]
+    model_path = f"checkpoints/{model_name}"
+    model = load_model(model_path, feature_dim)
+    models.append(model_name)
+    
+setup(models_info[0])
 
 with gr.Blocks() as demo:
     # title