Update app.py

app.py

@@ -1,320 +1,320 @@
-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)
+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 'Improving Fine-Grained Image Classification Using Diffusion-Based Generated Synthetic Images' 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)