Zentochi/Demo-Ensemble-Learning

1import gradio as gr2import torch3import torch.nn as nn4import torch.nn.functional as F5from torchvision.transforms import v26from torchvision.models import resnet34, vgg16, mobilenet_v2, ResNet34_Weights, VGG16_Weights, MobileNet_V2_Weights7from PIL import Image8import os9import numpy as np10from collections import Counter11from localization import i18n12import spaces13 14device = torch.device("cpu")15CLASS_NAME = ['healthy_leaf', 'others', 'ringspot_leaf', 'yellow_leaf']16BASE_LEARNER = ["VGG16", "ResNet34", "MobileNetV2"]17 18# Dropdown Gradio choice list19GRADIO_CHOICES = [20    "ResNet34",21    "VGG16",22    "MobileNetV2",23    "Ensemble Learning (Soft Voting)",24]25 26class CustomNet(nn.Module):27    def __init__(self, output_size=len(CLASS_NAME), architecture="VGG16"):28        super().__init__()29        self.architecture_name = architecture30 31        if self.architecture_name == "ResNet34":32            self.pretrain_model = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1)33        elif self.architecture_name == "VGG16":34            self.pretrain_model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1)35        elif self.architecture_name == "MobileNetV2":36            self.pretrain_model = mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1)37        else:38            raise ValueError(f"Unsupported {architecture} architecture")39 40        self.freeze_layers()41 42        if self.architecture_name == "ResNet34":43            self.pretrain_model.fc = nn.Sequential(44                nn.Linear(512, output_size)45            )46        elif self.architecture_name == "VGG16":47            self.pretrain_model.classifier = nn.Sequential(48                nn.Linear(25088, 4096),49                nn.ReLU(inplace=True),50                nn.Dropout(p=0.5),51                nn.Linear(4096, 4096),52                nn.ReLU(inplace=True),53                nn.Dropout(p=0.5),54                nn.Linear(4096, output_size)55            )56        elif self.architecture_name == "MobileNetV2":57            self.pretrain_model.classifier = nn.Sequential(58                nn.Dropout(p=0.2, inplace=False),59                nn.Linear(1280, output_size)60            )61 62    def forward(self, x):63        return self.pretrain_model(x)64 65    def freeze_layers(self):66        for param in self.pretrain_model.parameters():67            param.requires_grad = False68 69    def unfreeze_layers(self):70        for param in self.pretrain_model.parameters():71            param.requires_grad = True72 73global_loaded_models = {}74 75print("Memuat model-model ensemble...")76for arch_name in BASE_LEARNER:77    model = CustomNet(len(CLASS_NAME), architecture=arch_name)78    weights_path = f"./model-ensemble/best_model_{arch_name.lower()}_weight.pth"79 80    if os.path.exists(weights_path):81        try:82            state_dict = torch.load(weights_path, map_location=device)83            model.load_state_dict(state_dict)84            print(f"  Berhasil memuat weight untuk {arch_name} dari {weights_path}")85            model.to(device)86            model.eval()87            global_loaded_models[arch_name.lower()] = model88        except Exception as e:89            print(f"  [ERROR] Gagal memuat state_dict untuk {arch_name} dari {weights_path}: {e}")90    else:91        print(f"  [WARN] Custom weight tidak ditemukan untuk {arch_name} di {weights_path}. Model ini tidak akan digunakan dalam ensemble atau sebagai model tunggal.")92 93if not global_loaded_models:94    raise RuntimeError("Tidak ada model yang berhasil dimuat. Pastikan path weight.pth sudah benar dan file tersedia.")95 96def preprocess_image(image: Image.Image, target_device: torch.device):97    CROP_SIZE = (224, 224)98    image_transforms = v2.Compose([99        v2.Resize(CROP_SIZE),100        #v2.CenterCrop(CROP_SIZE),101        v2.ToImage(),102        v2.ToDtype(torch.float32, scale=True),103        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])104    ])105    image = image_transforms(image)106    image = image.unsqueeze(0)107    return image.to(target_device)108 109 110# Add ZeroGPU Configuration Annotation111@spaces.GPU112def model_predict(image_input: Image.Image, mode: str, selected_option: str) -> dict:113    if image_input is None:114        return {class_name: 0.0 for class_name in CLASS_NAME}115    # Dynamically target the GPU context assigned by ZeroGPU116    run_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")117 118    if mode == "User":119        selected_option = "Ensemble Learning (Soft Voting)"120 121    image_input = image_input.convert("RGB")122    processed_image = preprocess_image(image_input, run_device)123 124    final_probabilities = None125 126    with torch.no_grad():127        if selected_option in BASE_LEARNER:128            model_name_lower = selected_option.lower()129            if model_name_lower not in global_loaded_models:130                print(f"Error: Model {selected_option} tidak dimuat.")131                return {class_name: 0.0 for class_name in CLASS_NAME}132 133            model = global_loaded_models[model_name_lower].to(run_device)134            outputs = model(processed_image)135            final_probabilities = F.softmax(outputs, dim=1).cpu().numpy()[0]136 137        elif selected_option == "Ensemble Learning (Soft Voting)":138            all_model_raw_outputs = []139            for model_name in BASE_LEARNER:140                if model_name.lower() in global_loaded_models:141                    model = global_loaded_models[model_name.lower()].to(run_device)142                    outputs = model(processed_image)143                    all_model_raw_outputs.append(outputs.cpu().numpy())144                else:145                    print(f"Peringatan: Model {model_name} tidak dimuat atau gagal dimuat, dilewati dari ensemble.")146 147            if not all_model_raw_outputs:148                print("Error: Tidak ada model yang tersedia untuk ensemble (soft voting).")149                return {class_name: 0.0 for class_name in CLASS_NAME}150 151            all_model_raw_outputs_np = np.stack(all_model_raw_outputs, axis=0)152            all_model_raw_outputs_np = all_model_raw_outputs_np.transpose(1, 2, 0)153            154            all_model_probabilities = F.softmax(torch.from_numpy(all_model_raw_outputs_np), dim=1)155            final_probabilities = torch.mean(all_model_probabilities, dim=2).cpu().numpy()[0]156 157        else:158            raise ValueError(f"Opsi yang tidak valid: {selected_option}")159 160    prediction_dict = {CLASS_NAME[i]: float(final_probabilities[i]) for i in range(len(CLASS_NAME))}161    return prediction_dict162 163def update_model_dropdown_visibility(mode_choice):164    if mode_choice == "User":165        # Sembunyikan dropdown model dan kembalikan nilai default Soft Voting166        return gr.Dropdown(visible=False, value="Ensemble Learning (Soft Voting)")167    else: # mode_choice == "Tester"168        # Tampilkan dropdown model dan kembalikan nilai default Soft Voting169        return gr.Dropdown(visible=True, value="Ensemble Learning (Soft Voting)")170 171# Gradio Interface Code172with gr.Blocks() as demo:173    #title = gr.Markdown("<center><h1>Klasifikasi Citra Penyakit Daun Pepaya dengan <em>Ensemble Learning</em></h1></center>")174    #description = gr.Markdown("""<h3><center>Demo Tugas Akhir (TA) berjudul "Klasifikasi Citra Penyakit Daun Pepaya dengan <em>Ensemble Learning</em>". Prediksi citra menggunakan gabungan dari <em>Convolutional Neural Network</em> (CNN). Kelas prediksi yang tersedia adalah <em>Healthy leaf, Ringspot leaf, Yellow Leaf, </em>dan <em>Others</em>. Dibangun dengan <em>Gradio</em> dan <em>PyTorch</em>🔥</center></h3>""")175    gr.Markdown(i18n("title"))176    gr.Markdown(i18n("description"))177 178    with gr.Row(variant="panel", equal_height=True):179        image_input = gr.Image(height=512, width=512, type='pil', label=i18n("input_image"), sources=['upload', 'clipboard'])180        output_prediction_component = gr.Label(num_top_classes=len(CLASS_NAME), label=i18n("result"))181 182 183    example_list = []184    if os.path.exists("examples"):185        example_list = [["examples/" + example] for example in os.listdir("examples") if example.lower().endswith(('.png', '.jpg', '.jpeg'))]186 187    if example_list:188        gr.Examples(189            examples=example_list,190            inputs=[image_input],191            label=i18n("image_example"),192            cache_examples=False193        )194    else:195        gr.Markdown(i18n("missing_image"))196 197    with gr.Column(variant="panel"):198        mode_choice = gr.Dropdown(199            choices=["User", "Tester"],200            value="User", # Default mode201            label=i18n("mode_selector"),202            info=i18n("mode_description"),203            allow_custom_value=False,204            filterable=False,205        )206 207        model_selection_choice = gr.Dropdown(208            choices=GRADIO_CHOICES,209            value="Ensemble Learning (Soft Voting)", # Default Choice210            label=i18n("choose_method"),211            info=i18n("method_description"),212            allow_custom_value=False,213            filterable=False,214            visible=False # Initialize hidden215        )216        with gr.Row(variant="panel"):217            run_button = gr.Button(i18n("predict_button"), variant='primary')218            clear_button = gr.Button(i18n("clear_button"), variant='secondary')219 220    mode_choice.change(221        fn=update_model_dropdown_visibility,222        inputs=[mode_choice],223        outputs=[model_selection_choice],224        api_name="update_dropdown_visibility"225    )226 227    run_button.click(228        inputs=[image_input, mode_choice, model_selection_choice],229        outputs=output_prediction_component,230        fn=model_predict231    )232 233    clear_button.click(234        outputs=[image_input, output_prediction_component, mode_choice, model_selection_choice],235        fn=lambda: (None, {None: None}, "User", gr.Dropdown(visible=False, value="Ensemble Learning (Soft Voting)"))236    )237 238if __name__ == "__main__":239    demo.launch(i18n=i18n)240