Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from torchvision import transforms | |
| from PIL import Image | |
| import json, random, os | |
| # --- 1️⃣ Load models --- | |
| model_paths = { | |
| "Classifier Only (M1)": "M1.pt", | |
| "Classifier + L1-4 (M3)": "M3.pt", | |
| "Classifier + L1-4 (M4)": "M4.pt" | |
| } | |
| models = {name: torch.jit.load(path, map_location='cpu').eval() | |
| for name, path in model_paths.items()} | |
| # --- 2️⃣ Load label map --- | |
| with open("label_map.json") as f: | |
| label_map = json.load(f) | |
| classes = [label_map[k] for k in sorted(label_map, key=int)] | |
| # --- 3️⃣ Load test data --- | |
| TEST_DIR = "test" | |
| TEST_FILE = "test.txt" | |
| with open(TEST_FILE, "r") as f: | |
| test_entries = [line.strip().split() for line in f.readlines()] | |
| test_data = [(os.path.join(TEST_DIR, fn), int(cls)) for fn, cls in test_entries] | |
| # --- 4️⃣ Preprocessing --- | |
| preprocess = transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| ]) | |
| # --- 5️⃣ Helper functions --- | |
| def get_random_test_image(): | |
| """Pick a random test image and its label.""" | |
| img_path, cls_id = random.choice(test_data) | |
| true_label = label_map[str(cls_id)] | |
| img = Image.open(img_path).convert("RGB") | |
| return img, f"Ground Truth: {true_label}" | |
| def classify_bird(img, model_name, mode): | |
| """Classify based on current input mode.""" | |
| if img is None: | |
| return {} | |
| model = models[model_name] | |
| x = preprocess(img).unsqueeze(0) | |
| with torch.no_grad(): | |
| logits = model(x) | |
| probs = torch.softmax(logits, dim=1)[0] | |
| top5_prob, top5_idx = torch.topk(probs, 5) | |
| return {classes[i]: float(top5_prob[j]) for j, i in enumerate(top5_idx)} | |
| # --- 6️⃣ Interface layout --- | |
| with gr.Blocks(theme=gr.themes.Soft(), title="CUB-200 Bird Classifier") as demo: | |
| gr.Markdown(""" | |
| # 🐦 CUB-200 Bird Classifier | |
| Identify bird species using fine-tuned ResNet101 models. | |
| Choose input mode and model, then classify! | |
| """) | |
| # --- Input controls --- | |
| with gr.Row(): | |
| input_mode = gr.Radio( | |
| choices=["Upload Image", "Use Random Test Image"], | |
| value="Upload Image", | |
| label="Input Mode", | |
| interactive=True | |
| ) | |
| model_choice = gr.Dropdown( | |
| choices=list(models.keys()), | |
| value="Classifier + L1-4 (M4)", | |
| label="Select Model" | |
| ) | |
| gr.Markdown("---") | |
| # --- Main visualization area --- | |
| with gr.Row(equal_height=True): | |
| with gr.Column(): | |
| image_display = gr.Image(type="pil", label="Bird Image", interactive=True) | |
| regenerate_btn = gr.Button("🔁 Regenerate Random Image", visible=False) | |
| ground_truth = gr.Textbox(label="Ground Truth", interactive=False, visible=False) | |
| predictions = gr.Label(num_top_classes=5, label="Top-5 Predictions") | |
| # --- Bottom control --- | |
| classify_btn = gr.Button("🔍 Classify") | |
| # --- Logic --- | |
| def handle_mode_change(mode): | |
| """Handle mode switch and optionally load a random image.""" | |
| if mode == "Use Random Test Image": | |
| img, label = get_random_test_image() | |
| return ( | |
| gr.update(value=img, interactive=False), | |
| gr.update(visible=True), | |
| gr.update(value=label, visible=True) | |
| ) | |
| else: | |
| return ( | |
| gr.update(value=None, interactive=True), | |
| gr.update(visible=False), | |
| gr.update(value="", visible=False) | |
| ) | |
| def regenerate_random_image(): | |
| img, label = get_random_test_image() | |
| return gr.update(value=img), gr.update(value=label) | |
| input_mode.change( | |
| fn=handle_mode_change, | |
| inputs=input_mode, | |
| outputs=[image_display, regenerate_btn, ground_truth] | |
| ) | |
| regenerate_btn.click( | |
| fn=regenerate_random_image, | |
| inputs=None, | |
| outputs=[image_display, ground_truth] | |
| ) | |
| classify_btn.click( | |
| fn=classify_bird, | |
| inputs=[image_display, model_choice, input_mode], | |
| outputs=predictions | |
| ) | |
| demo.launch() | |