import gradio as gr import torch import torch.nn.functional as F from torch.utils.data import DataLoader import matplotlib.pyplot as plt from model_module import AutoencoderModule 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 import utils import spaces def load_model(model_path, feature_dim): model = AutoencoderModule(feature_dim=feature_dim) state_dict = torch.load(model_path) if "state_dict" in state_dict: model.load_state_dict(state_dict['state_dict']) model.eval() else: # state_dict のキーを修正 new_state_dict = {} for key in state_dict: new_key = "model." + key new_state_dict[new_key] = state_dict[key] model.load_state_dict(new_state_dict) model.eval() model.to(device) print(f"{model_path} loaded successfully.") return model def load_data(img_dir="resources/trainB/", image_size=112, batch_size=256): filenames = load_filenames(img_dir) train_X = filenames[:1000] train_ds = MyDataset(train_X, img_dir=img_dir, img_size=image_size) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0) iterator = iter(train_loader) x, _, _ = next(iterator) x = x.to(device) x = x[:,0].to(device) print("Data loaded successfully.") return x 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') image_points_ds = ImageKeypointDataset(train_X, keypoints, img_dir='resources/trainB/', img_size=image_size) image_points_loader = DataLoader(image_points_ds, batch_size=batch_size, shuffle=False) iterator = iter(image_points_loader) test_imgs, points = next(iterator) test_imgs = test_imgs.to(device) points = points.to(device)*(image_size) print("Keypoints loaded successfully.") return test_imgs, points image_size = 112 batch_size = 32 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") models_info = [ {"name": "autoencoder-epoch=49-train_loss=1.01.ckpt", "feature_dim": 64}, {"name": "autoencoder-epoch=29-train_loss=1.01.ckpt", "feature_dim": 64}, {"name": "autoencoder-epoch=09-train_loss=1.00.ckpt", "feature_dim": 64}, {"name": "ae_model_tf_2024-03-05_00-35-21.pth", "feature_dim": 32}, ] models = [] for model_info in models_info: model_name = model_info["name"] feature_dim = model_info["feature_dim"] model_path = f"checkpoints/{model_name}" models.append(load_model(model_path, feature_dim)) x = load_data() test_imgs, points = load_keypoints() mean_vector_list = [] model_index = 0 # ヒートマップの生成関数 @spaces.GPU 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: 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) 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 # アップロード画像の前処理 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).to(device) target_feature_map, _ = models[model_index](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 with gr.Blocks() as demo: # title gr.Markdown("# TripletGeoEncoder Feature Map Visualization") # description gr.Markdown("This demo visualizes the feature maps of a TripletGeoEncoder trained on the CelebA dataset using self-supervised learning without annotations from only 1000 images. " "The feature maps are visualized as heatmaps, where the source map shows the distance of each pixel in the source image to the selected pixel, and the target map shows the distance of each pixel in the target image to the selected pixel. " "The blended source and target images show the source and target images with the source and target maps overlaid, respectively. " "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) 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"), input_image ], outputs=output_plot, live=True, flagging_mode="never" ) # examples gr.Markdown("# Examples") gr.Examples( examples=[ ["resources/examples/2488.jpg"], ["resources/examples/2899.jpg"] ], inputs=[input_image], ) demo.launch()