app.py

import random

import gradio as gr
import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch import nn
from torchvision.models import mobilenet_v2, resnet18
from torchvision.transforms.functional import InterpolationMode

datasets_n_classes = {
    "Imagenette": 10,
    "Imagewoof": 10,
    "Stanford_dogs": 120,
}

datasets_model_types = {
    "Imagenette": [
        "base_200",
        "base_200+100",
        "synthetic_200",
        "augment_noisy_200",
        "augment_noisy_200+100",
        "augment_clean_200",
    ],
    "Imagewoof": [
        "base_200",
        "base_200+100",
        "synthetic_200",
        "augment_noisy_200",
        "augment_noisy_200+100",
        "augment_clean_200",
    ],
    "Stanford_dogs": [
        "base_200",
        "base_200+100",
        "synthetic_200",
        "augment_noisy_200",
        "augment_noisy_200+100",
    ],
}

model_arch = ["resnet18", "mobilenet_v2"]

list_200 = [
    "Original",
    "Synthetic",
    "Original + Synthetic (Noisy)",
    "Original + Synthetic (Clean)",
]

list_200_100 = ["Base+100", "AugmentNoisy+100"]

methods_map = {
    "200 Epochs": list_200,
    "200 Epochs on Original + 100": list_200_100,
}

label_map = dict()
label_map["Imagenette (10 classes)"] = "Imagenette"
label_map["Imagewoof (10 classes)"] = "Imagewoof"
label_map["Stanford Dogs (120 classes)"] = "Stanford_dogs"
label_map["ResNet-18"] = "resnet18"
label_map["MobileNetV2"] = "mobilenet_v2"
label_map["200 Epochs"] = "200"
label_map["200 Epochs on Original + 100"] = "200+100"
label_map["Original"] = "base"
label_map["Synthetic"] = "synthetic"
label_map["Original + Synthetic (Noisy)"] = "augment_noisy"
label_map["Original + Synthetic (Clean)"] = "augment_clean"
label_map["Base+100"] = "base"
label_map["AugmentNoisy+100"] = "augment_noisy"

dataset_models = dict()
for dataset, n_classes in datasets_n_classes.items():
    models = dict()
    for model_type in datasets_model_types[dataset]:
        for arch in model_arch:
            if arch == "resnet18":
                model = resnet18(weights=None, num_classes=n_classes)
                models[f"{arch}_{model_type}"] = (
                    model,
                    f"./models/{arch}/{dataset}/{dataset}_{model_type}.pth",
                )
            elif arch == "mobilenet_v2":
                model = mobilenet_v2(weights=None, num_classes=n_classes)
                models[f"{arch}_{model_type}"] = (
                    model,
                    f"./models/{arch}/{dataset}/{dataset}_{model_type}.pth",
                )
            else:
                raise ValueError(f"Model architecture unavailable: {arch}")
    dataset_models[dataset] = models


def get_random_image(dataset, label_map=label_map) -> Image:
    dataset_root = f"./data/{label_map[dataset]}/val"
    dataset_img = torchvision.datasets.ImageFolder(
        dataset_root,
        transforms.Compose([transforms.PILToTensor()]),
    )
    random_idx = random.randint(0, len(dataset_img) - 1)
    image, _ = dataset_img[random_idx]
    image = transforms.ToPILImage()(image)
    image = image.resize(
        (256, 256),
    )
    return image


def load_model(model_dict, model_name: str) -> nn.Module:
    model_name_lower = model_name.lower()
    if model_name_lower in model_dict:
        model = model_dict[model_name_lower][0]
        model_path = model_dict[model_name_lower][1]
        if torch.cuda.is_available():
            checkpoint = torch.load(model_path)
        else:
            checkpoint = torch.load(model_path, map_location="cpu")
        if "setup" in checkpoint:
            if checkpoint["setup"]["distributed"]:
                torch.nn.modules.utils.consume_prefix_in_state_dict_if_present(
                    checkpoint["model"], "module."
                )
            model.load_state_dict(checkpoint["model"])
        else:
            model.load_state_dict(checkpoint)
        return model
    else:
        raise ValueError(
            f"Model {model_name} is not available for image prediction. Please choose from {[name.capitalize() for name in model_dict.keys()]}."
        )


def postprocess_default(labels, output) -> dict:
    probabilities = nn.functional.softmax(output[0], dim=0)
    top_prob, top_catid = torch.topk(probabilities, 5)
    confidences = {
        labels[top_catid.tolist()[i]]: top_prob.tolist()[i]
        for i in range(top_prob.shape[0])
    }
    return confidences


def classify(
    input_image: Image,
    dataset_type: str,
    arch_type: str,
    methods: str,
    training_ds: str,
    dataset_models=dataset_models,
    label_map=label_map,
) -> dict:
    for i in [dataset_type, arch_type, methods, training_ds]:
        if i is None:
            raise ValueError("Please select all options.")
    dataset_type = label_map[dataset_type]
    arch_type = label_map[arch_type]
    methods = label_map[methods]
    training_ds = label_map[training_ds]
    preprocess_input = transforms.Compose(
        [
            transforms.Resize(
                256,
                interpolation=InterpolationMode.BILINEAR,
                antialias=True,
            ),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    if input_image is None:
        raise ValueError("No image was provided.")
    input_tensor: torch.Tensor = preprocess_input(input_image)
    input_batch = input_tensor.unsqueeze(0)
    model = load_model(
        dataset_models[dataset_type], f"{arch_type}_{training_ds}_{methods}"
    )

    if torch.cuda.is_available():
        input_batch = input_batch.to("cuda")
        model.to("cuda")

    model.eval()
    with torch.inference_mode():
        output: torch.Tensor = model(input_batch)
    with open(f"./data/{dataset_type}.txt", "r") as f:
        labels = {i: line.strip() for i, line in enumerate(f.readlines())}
    return postprocess_default(labels, output)


def update_methods(method, ds_type):
    if ds_type == "Stanford Dogs (120 classes)" and method == "200 Epochs":
        methods = list_200[:-1]
    else:
        methods = methods_map[method]
    return gr.update(choices=methods, value=None)


def downloadModel(
    dataset_type, arch_type, methods, training_ds, dataset_models=dataset_models
):
    for i in [dataset_type, arch_type, methods, training_ds]:
        if i is None:
            return gr.update(label="Select Model", value=None)
    dataset_type = label_map[dataset_type]
    arch_type = label_map[arch_type]
    methods = label_map[methods]
    training_ds = label_map[training_ds]
    if f"{arch_type}_{training_ds}_{methods}" not in dataset_models[dataset_type]:
        return gr.update(label="Select Model", value=None)
    model_path = dataset_models[dataset_type][f"{arch_type}_{training_ds}_{methods}"][1]
    return gr.update(
        label=f"Download Model: '{dataset_type}_{arch_type}_{training_ds}_{methods}'",
        value=model_path,
    )


if __name__ == "__main__":
    with gr.Blocks(title="Generative Augmented Image Classifiers") as demo:
        gr.Markdown(
            """
# Generative Augmented Image Classifiers
Main GitHub Repo: [Generative Data Augmentation](https://github.com/zhulinchng/generative-data-augmentation) | Generative Data Augmentation Demo: [Generative Data Augmented](https://huggingface.co/spaces/czl/generative-data-augmentation-demo).
"""
        )
        with gr.Row():
            with gr.Column():
                dataset_type = gr.Radio(
                    choices=[
                        "Imagenette (10 classes)",
                        "Imagewoof (10 classes)",
                        "Stanford Dogs (120 classes)",
                    ],
                    label="Dataset",
                    value="Imagenette (10 classes)",
                )
                arch_type = gr.Radio(
                    choices=["ResNet-18", "MobileNetV2"],
                    label="Model Architecture",
                    value="ResNet-18",
                    interactive=True,
                )
                methods = gr.Radio(
                    label="Methods",
                    choices=["200 Epochs", "200 Epochs on Original + 100"],
                    interactive=True,
                    value="200 Epochs",
                )
                training_ds = gr.Radio(
                    label="Training Dataset",
                    choices=methods_map["200 Epochs"],
                    interactive=True,
                    value="Original",
                )
                dataset_type.change(
                    fn=update_methods,
                    inputs=[methods, dataset_type],
                    outputs=[training_ds],
                )
                methods.change(
                    fn=update_methods,
                    inputs=[methods, dataset_type],
                    outputs=[training_ds],
                )
                random_image_output = gr.Image(type="pil", label="Image to Classify")
                with gr.Row():
                    generate_button = gr.Button("Sample Random Image")
                    classify_button_random = gr.Button("Classify")
            with gr.Column():
                output_label_random = gr.Label(num_top_classes=5)
                download_model = gr.DownloadButton(
                    label=f"Download Model: '{label_map[dataset_type.value]}_{label_map[arch_type.value]}_{label_map[training_ds.value]}_{label_map[methods.value]}'",
                    value=dataset_models[label_map[dataset_type.value]][
                        f"{label_map[arch_type.value]}_{label_map[training_ds.value]}_{label_map[methods.value]}"
                    ][1],
                )
                dataset_type.change(
                    fn=downloadModel,
                    inputs=[dataset_type, arch_type, methods, training_ds],
                    outputs=[download_model],
                )
                arch_type.change(
                    fn=downloadModel,
                    inputs=[dataset_type, arch_type, methods, training_ds],
                    outputs=[download_model],
                )
                methods.change(
                    fn=downloadModel,
                    inputs=[dataset_type, arch_type, methods, training_ds],
                    outputs=[download_model],
                )
                training_ds.change(
                    fn=downloadModel,
                    inputs=[dataset_type, arch_type, methods, training_ds],
                    outputs=[download_model],
                )
                gr.Markdown(
                    """
This demo showcases the performance of image classifiers trained on various datasets as part of the project 'Investigating the Effectiveness of Generative Diffusion Models in Synthesizing Images for Data Augmentation in Image Classification' dissertation.

View the models and files used in this demo [here](https://huggingface.co/spaces/czl/generative-augmented-classifiers/tree/main).

Usage Instructions & Documentation [here](https://huggingface.co/spaces/czl/generative-augmented-classifiers/blob/main/README.md).
                """
                )

        generate_button.click(
            get_random_image,
            inputs=[dataset_type],
            outputs=random_image_output,
        )
        classify_button_random.click(
            classify,
            inputs=[random_image_output, dataset_type, arch_type, methods, training_ds],
            outputs=output_label_random,
        )
    demo.launch(show_error=True)