app.py

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 dataset
from dataset import MyDataset, ImageKeypointDataset, load_filenames, load_keypoints
import utils

try:
    import spaces
except ImportError:
    print("Spaces is not installed.")

image_size = 112
batch_size = 32


# モデルとデータの読み込み
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)
    
    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()
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    print("Model loaded successfully.")
    return model, device

def load_data(device, 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(device, 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

# ヒートマップの生成関数
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)
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)
        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])


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")

    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)
    with gr.Row():
        with gr.Column():
            with gr.Row():
                model_name = gr.Dropdown(
                    choices=[str(model) for model in models],
                    container=False
                )
            load_button = gr.Button("Load Model")
            load_button.click(setup, inputs=[model_name, input_image], outputs=[output_plot])
        with gr.Row():
            pass    

    inference = gr.Interface(
        get_heatmaps,
        inputs=[
            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],
    )
            
    # JavaScriptコードをロード
    demo.launch()