gpu

app.py

@@ -13,16 +13,12 @@ import dataset
 from dataset import MyDataset, ImageKeypointDataset, load_filenames, load_keypoints
 import utils
 
-try:
-    import spaces
-except ImportError:
-    print("Spaces is not installed.")
+import spaces
 
 image_size = 112
 batch_size = 32
 
-
-# モデルとデータの読み込み
+@spaces.GPU
 def load_model(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)
@@ -78,98 +74,51 @@ def load_keypoints(device, img_dir="resources/trainB/", image_size=112, batch_si
     return test_imgs, points
 
 # ヒートマップの生成関数
-try:
-    @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)
-
-        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
-
-except:
-    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
+@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)
+
+    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