gpu

app.py

@@ -125,31 +125,6 @@ try:
         plt.close(fig)
         return fig
 
-    def setup(model_dict, input_image=None):
-        global model, device, 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, device = load_model(model_path, feature_dim)
-        x = load_data(device)
-        test_imgs, points = load_keypoints(device)
-        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},
-            {"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])
 except:
     def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
         if type(uploaded_image) == str:
@@ -196,32 +171,31 @@ except:
         plt.close(fig)
         return fig
 
-    def setup(model_dict, input_image=None):
-        global model, device, 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, device = load_model(model_path, feature_dim)
-        x = load_data(device)
-        test_imgs, points = load_keypoints(device)
-        feature_map, _ = model(test_imgs)
-        mean_vector_list = utils.get_mean_vector(feature_map, points)
+def setup(model_dict, input_image=None):
+    global model, device, 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, device = load_model(model_path, feature_dim)
+    x = load_data(device)
+    test_imgs, points = load_keypoints(device)
+    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
         
-        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},
-            {"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 = [{"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])
 
 with gr.Blocks() as demo:
     # title