gpu

app.py

@@ -81,81 +81,146 @@ def load_keypoints(device, img_dir="resources/trainB/", image_size=112, batch_si
 try:
     @spaces.GPU
     def get_heatmaps(source_num, x_coords, y_coords, uploaded_image):
-        return _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)
+
+        dec5, _ = model(x)
+        feature_map = dec5
+        # アップロード画像の前処理
+        if uploaded_image is not None:
+            uploaded_image = utils.preprocess_uploaded_image(uploaded_image['composite'], image_size)
+        else:
+            uploaded_image = torch.zeros(1, 3, image_size, image_size, device=device)
+        target_feature_map, _ = model(uploaded_image)
+        img = torch.cat((x, uploaded_image))
+        feature_map = torch.cat((feature_map, target_feature_map))
+
+        source_map, target_map, blended_source, blended_target = utils.get_heatmaps(img, feature_map, source_num, x_coords, y_coords, uploaded_image)
+        keypoint_maps, blended_tensors = utils.get_keypoint_heatmaps(target_feature_map, mean_vector_list, points.size(1), uploaded_image)
+
+        # Matplotlibでプロットして画像として保存
+        fig, axs = plt.subplots(2, 3, figsize=(10, 6))
+        axs[0, 0].imshow(source_map, cmap='hot')
+        axs[0, 0].set_title("Source Map")
+        axs[0, 1].imshow(target_map, cmap='hot')
+        axs[0, 1].set_title("Target Map")
+        axs[0, 2].imshow(keypoint_maps[0], cmap='hot')
+        axs[0, 2].set_title("Keypoint Map")
+        axs[1, 0].imshow(blended_source.permute(1, 2, 0))
+        axs[1, 0].set_title("Blended Source")
+        axs[1, 1].imshow(blended_target.permute(1, 2, 0))
+        axs[1, 1].set_title("Blended Target")
+        axs[1, 2].imshow(blended_tensors[0].permute(1, 2, 0))
+        axs[1, 2].set_title("Blended Keypoint")
+        for ax in axs.flat:
+            ax.axis('off')
+        
+        plt.tight_layout()
+        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):
-        return _get_heatmaps(source_num, x_coords, y_coords, uploaded_image)
-    
-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)
-
-    dec5, _ = model(x)
-    feature_map = dec5
-    # アップロード画像の前処理
-    if uploaded_image is not None:
-        uploaded_image = utils.preprocess_uploaded_image(uploaded_image['composite'], image_size)
-    else:
-        uploaded_image = torch.zeros(1, 3, image_size, image_size, device=device)
-    target_feature_map, _ = model(uploaded_image)
-    img = torch.cat((x, uploaded_image))
-    feature_map = torch.cat((feature_map, target_feature_map))
-
-    source_map, target_map, blended_source, blended_target = utils.get_heatmaps(img, feature_map, source_num, x_coords, y_coords, uploaded_image)
-    keypoint_maps, blended_tensors = utils.get_keypoint_heatmaps(target_feature_map, mean_vector_list, points.size(1), uploaded_image)
-
-    # Matplotlibでプロットして画像として保存
-    fig, axs = plt.subplots(2, 3, figsize=(10, 6))
-    axs[0, 0].imshow(source_map, cmap='hot')
-    axs[0, 0].set_title("Source Map")
-    axs[0, 1].imshow(target_map, cmap='hot')
-    axs[0, 1].set_title("Target Map")
-    axs[0, 2].imshow(keypoint_maps[0], cmap='hot')
-    axs[0, 2].set_title("Keypoint Map")
-    axs[1, 0].imshow(blended_source.permute(1, 2, 0))
-    axs[1, 0].set_title("Blended Source")
-    axs[1, 1].imshow(blended_target.permute(1, 2, 0))
-    axs[1, 1].set_title("Blended Target")
-    axs[1, 2].imshow(blended_tensors[0].permute(1, 2, 0))
-    axs[1, 2].set_title("Blended Keypoint")
-    for ax in axs.flat:
-        ax.axis('off')
-    
-    plt.tight_layout()
-    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)
+        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)
+
+        dec5, _ = model(x)
+        feature_map = dec5
+        # アップロード画像の前処理
+        if uploaded_image is not None:
+            uploaded_image = utils.preprocess_uploaded_image(uploaded_image['composite'], image_size)
+        else:
+            uploaded_image = torch.zeros(1, 3, image_size, image_size, device=device)
+        target_feature_map, _ = model(uploaded_image)
+        img = torch.cat((x, uploaded_image))
+        feature_map = torch.cat((feature_map, target_feature_map))
+
+        source_map, target_map, blended_source, blended_target = utils.get_heatmaps(img, feature_map, source_num, x_coords, y_coords, uploaded_image)
+        keypoint_maps, blended_tensors = utils.get_keypoint_heatmaps(target_feature_map, mean_vector_list, points.size(1), uploaded_image)
+
+        # Matplotlibでプロットして画像として保存
+        fig, axs = plt.subplots(2, 3, figsize=(10, 6))
+        axs[0, 0].imshow(source_map, cmap='hot')
+        axs[0, 0].set_title("Source Map")
+        axs[0, 1].imshow(target_map, cmap='hot')
+        axs[0, 1].set_title("Target Map")
+        axs[0, 2].imshow(keypoint_maps[0], cmap='hot')
+        axs[0, 2].set_title("Keypoint Map")
+        axs[1, 0].imshow(blended_source.permute(1, 2, 0))
+        axs[1, 0].set_title("Blended Source")
+        axs[1, 1].imshow(blended_target.permute(1, 2, 0))
+        axs[1, 1].set_title("Blended Target")
+        axs[1, 2].imshow(blended_tensors[0].permute(1, 2, 0))
+        axs[1, 2].set_title("Blended Keypoint")
+        for ax in axs.flat:
+            ax.axis('off')
+        
+        plt.tight_layout()
+        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}]
+    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])
+    setup(models[0])
 
 
 with gr.Blocks() as demo: