| 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] | |
| checkpoint = torch.load(model_path) | |
| 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 | |
| This demo showcases the performance of image classifiers trained on various datasets. | |
| """ | |
| ) | |
| 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], | |
| ) | |
| generate_button = gr.Button("Sample Random Image") | |
| random_image_output = gr.Image( | |
| type="pil", label="Random Image from Validation Set" | |
| ) | |
| 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], | |
| ) | |
| 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) | |