Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import timm | |
| import torchvision.transforms as T | |
| from PIL import Image | |
| import numpy as np | |
| import pandas as pd | |
| import json, os | |
| # ------------------------------ | |
| # CONFIGURATION | |
| # ------------------------------ | |
| DEVICE = "cpu" | |
| NUM_CLASSES = 200 | |
| MODEL_PATHS = { | |
| "EfficientNet-B3": "models/BestEfficientNetB3_model.pth", | |
| "EfficientNetV2-S (in21k)": "models/BestEfficientNetV2S_model.pth" | |
| } | |
| MODEL_NAMES = { | |
| "EfficientNet-B3": "tf_efficientnet_b3_ns", | |
| "EfficientNetV2-S (in21k)": "tf_efficientnetv2_s.in21k" | |
| } | |
| # ------------------------------ | |
| # LOAD CLASS NAMES (supports JSON or TXT) | |
| # ------------------------------ | |
| CLASS_NAMES = [] | |
| if os.path.exists("class_names.json"): | |
| try: | |
| with open("class_names.json", "r") as f: | |
| data = json.load(f) | |
| if isinstance(data, dict): | |
| # some JSONs store indices as keys | |
| CLASS_NAMES = [v for _, v in sorted(data.items(), key=lambda x: int(x[0]))] | |
| elif isinstance(data, list): | |
| CLASS_NAMES = data | |
| print(f"✅ Loaded {len(CLASS_NAMES)} class names from class_names.json") | |
| except Exception as e: | |
| print(f"⚠️ Error reading class_names.json: {e}") | |
| elif os.path.exists("class_names.txt"): | |
| with open("class_names.txt", "r") as f: | |
| CLASS_NAMES = [line.strip() for line in f.readlines()] | |
| print(f"✅ Loaded {len(CLASS_NAMES)} class names from class_names.txt") | |
| else: | |
| CLASS_NAMES = [f"Class {i}" for i in range(NUM_CLASSES)] | |
| print("⚠️ No class_names file found — using generic labels") | |
| # ------------------------------ | |
| # MODEL LOADING + CACHING | |
| # ------------------------------ | |
| def load_model(model_name): | |
| model = timm.create_model(MODEL_NAMES[model_name], | |
| pretrained=False, | |
| num_classes=NUM_CLASSES) | |
| model.load_state_dict(torch.load(MODEL_PATHS[model_name], | |
| map_location=DEVICE)) | |
| model.to(DEVICE).eval() | |
| return model | |
| _loaded_models = {} | |
| def get_model(model_name): | |
| if model_name not in _loaded_models: | |
| print(f"🔄 Loading {model_name} ...") | |
| _loaded_models[model_name] = load_model(model_name) | |
| return _loaded_models[model_name] | |
| # ------------------------------ | |
| # PREDICTION FUNCTION | |
| # ------------------------------ | |
| def predict(image, model_choice): | |
| if image is None: | |
| return pd.DataFrame(), "Please upload an image." | |
| model = get_model(model_choice) | |
| # Ensure RGB | |
| try: | |
| image = image.convert("RGB") | |
| except Exception as e: | |
| return pd.DataFrame(), f"❌ Error reading image: {e}" | |
| # Resize dynamically | |
| size = 300 if "b3" in model_choice.lower() else 384 | |
| preprocess = T.Compose([ | |
| T.Resize((size, size)), | |
| T.ToTensor(), | |
| T.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| ]) | |
| img_tensor = preprocess(image).unsqueeze(0).to(DEVICE) | |
| with torch.no_grad(): | |
| logits = model(img_tensor) | |
| probs = torch.nn.functional.softmax(logits, dim=1)[0] | |
| top5 = torch.topk(probs, 5) | |
| indices = top5.indices.cpu().numpy().tolist() | |
| scores = top5.values.cpu().numpy().tolist() | |
| # Map to readable names | |
| if CLASS_NAMES and len(CLASS_NAMES) > max(indices): | |
| names = [CLASS_NAMES[i] for i in indices] | |
| else: | |
| names = [f"Class {i}" for i in indices] | |
| # Make a DataFrame (this avoids 'undefined' issue) | |
| df = pd.DataFrame({ | |
| "Class": names, | |
| "Confidence (%)": [round(s * 100, 2) for s in scores] | |
| }) | |
| return df, f"✅ Prediction complete using {model_choice}" | |
| # ------------------------------ | |
| # GRADIO INTERFACE | |
| # ------------------------------ | |
| with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo: | |
| gr.Markdown( | |
| """ | |
| #Bird Species Classifier | |
| Upload a bird image or use your camera to identify its species. | |
| Supports multiple EfficientNet models and displays top-5 predictions. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| img_in = gr.Image(type="pil", | |
| label="Upload or Capture Image", | |
| sources=["upload", "webcam"]) | |
| model_sel = gr.Dropdown( | |
| choices=list(MODEL_PATHS.keys()), | |
| value="EfficientNetV2-S", | |
| label="Select Model") | |
| btn = gr.Button("🔍 Classify") | |
| with gr.Column(scale=1): | |
| df_out = gr.Dataframe(label="Top-5 Predictions", | |
| interactive=False) | |
| status = gr.Textbox(label="Status", interactive=False) | |
| btn.click(fn=predict, inputs=[img_in, model_sel], | |
| outputs=[df_out, status]) | |
| gr.Markdown( | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |