santosh3110/Ecoclassify-Wildlife_Classifier
0
1import os 2import json3import sys4import tempfile5import pandas as pd6import torch7import zipfile8import streamlit as st9from PIL import Image10from torchvision import models, transforms11import torch.nn as nn12from pathlib import Path13 14sys.path.append(os.path.join(os.path.dirname(__file__), "src"))15from ecoclassify import logger16from ecoclassify.config.configuration import ConfigurationManager17from ecoclassify.components.batch_inference import BatchInference18from ecoclassify.components.explanation_generator import ExplanationGenerator19from ecoclassify.components.fine_tuning import FineTuner20from ecoclassify.utils.common import load_json21 22# -------------------23# CONFIG24# -------------------25CONFIG = {26 "model_path": "artifacts/training/resnet_model.pth",27 "label_mapping_path": "artifacts/data_ingestion/extracted_data/label_mapping.json",28 "mean_std_path": "artifacts/training/logs/mean_std.json",29 "gradcam_target_layer": "layer4" # For ResNet50 last conv block30}31 32# -------------------33# HELPER FUNCTIONS34# -------------------35@st.cache_resource36def load_model_and_transforms():37 # Load label mapping38 with open(CONFIG["label_mapping_path"], "r") as f:39 label_map = json.load(f)40 if all(str(k).isdigit() for k in label_map.keys()):41 # Keys are numeric42 idx_to_class = {int(k): v for k, v in label_map.items()}43 else:44 # Keys are class names, values are numeric IDs45 idx_to_class = {v: k for k, v in label_map.items()}46 47 # Load mean/std48 stats = load_json(Path(CONFIG["mean_std_path"])) # FIXED: ensure Path type49 mean, std = stats["mean"], stats["std"]50 51 # Model52 model = models.resnet50(weights=None)53 model.fc = nn.Linear(model.fc.in_features, len(idx_to_class))54 model.load_state_dict(torch.load(CONFIG["model_path"], map_location="cpu"))55 model.eval()56 57 # Transform58 transform = transforms.Compose([59 transforms.Resize((224, 224)),60 transforms.ToTensor(),61 transforms.Normalize(mean=mean, std=std)62 ])63 return model, transform, idx_to_class, mean, std64 65def predict_image(model, transform, idx_to_class, image):66 img_t = transform(image).unsqueeze(0)67 with torch.no_grad():68 outputs = model(img_t)69 probs = torch.softmax(outputs, dim=1).squeeze().tolist()70 pred_idx = torch.argmax(outputs, dim=1).item()71 return idx_to_class[pred_idx], probs72 73# -------------------74# STREAMLIT UI75# -------------------76st.set_page_config(page_title="EcoClassify", layout="wide")77st.title("🦉 EcoClassify - Wildlife Image Classifier")78 79tabs = st.tabs(["ℹ️ About","📸 Inference", "📦 Batch Inference", "🔧 Fine-tuning"])80 81# -------------------82# TAB 1: ABOUT83# -------------------84with tabs[0]:85 st.header("🌍 About EcoClassify")86 87 st.markdown("""88 **Welcome to EcoClassify!** 🦉 89 Where **wildlife meets deep learning**. 90 91 **EcoClassify - Wildlife Image Classifier** was born out of the need to **help researchers, educators, and nature lovers** 92 quickly identify species captured in camera trap images — without needing to be a machine learning wizard.93 This project was developed as part of my internship at **[Euron](https://euron.one/)**, with heartfelt thanks to **Sudhanshu Kumar**, Director of Euron, for his guidance and support.94 95 ---96 ### 🚀 What it does97 1. **Classifies wildlife images** into 8 wildlife species — Antelope_Duiker, Bird, Civet_Genet, Hog, Leopard, Monkey_Prosimian, Rodent, 98 and yes… it can even recognize when there’s nothing there at all — just Blank 🙈.99 2. Uses **transfer learning with ResNet50** for strong, accurate predictions.100 3. **Explains predictions** with Grad-CAM heatmaps — so you can *see what the model sees*.101 4. Supports **batch processing** for test datasets.102 5. Lets you **fine-tune the model** on your own custom dataset, straight from the UI.103 104 ---105 ### 🛠 Under the Hood106 - **Frontend:** Streamlit — keeping it simple & interactive.107 - **Model:** PyTorch ResNet50, fine-tuned on a Wildlife Dataset collected from **drivendata.org**.108 - **Image Processing:** OpenCV + torchvision for augmentations & preprocessing.109 - **Explainability:** Grad-CAM visualizations via `torchcam`.110 - **Data Handling:** Pandas, NumPy.111 112 ---113 ### 🦜 Why it matters114 In the field of wildlife conservation, time matters. 115 Camera traps generate *thousands* of images, and manually sorting them 116 is both tedious and error-prone. 117 EcoClassify helps:118 - Researchers: Quickly analyze species distribution.119 - Educators: Teach students how AI models "think".120 - Wildlife enthusiasts: Get AI-powered insights on sightings.121 122 ---123 ### 📚 Dataset Reference124 The Pan African Programme: The Cultured Chimpanzee, Wild Chimpanzee Foundation, DrivenData. (2022). 125 *Conser-vision Practice Area: Image Classification.* 126 Retrieved [July 12, 2025] from 127 [https://www.drivendata.org/competitions/87/competition-image-classification-wildlife-conservation/](https://www.drivendata.org/competitions/87/competition-image-classification-wildlife-conservation/)128 129 ---130 131 ---132 **Made with ❤️ for Wildlife & AI by Santosh Kumar Guntupalli.** 133 _Let’s help protect the wild._134 """)135 136# -------------------137# TAB 2: INFERENCE138# -------------------139with tabs[1]:140 st.header("Single / Multiple Image Prediction with GradCAM")141 uploaded_files = st.file_uploader("Upload image(s)", type=["jpg", "jpeg", "png"], accept_multiple_files=True)142 143 if uploaded_files:144 model, transform, idx_to_class, mean, std = load_model_and_transforms()145 explainer_config = type("Config", (object,), {146 "model_weights": CONFIG["model_path"],147 "mean_std_path": CONFIG["mean_std_path"],148 "label_mapping_path": CONFIG["label_mapping_path"],149 "gradcam_target_layer": CONFIG["gradcam_target_layer"],150 "root_dir": "artifacts/streamlit_outputs"151 })()152 os.makedirs(explainer_config.root_dir, exist_ok=True)153 explainer = ExplanationGenerator(explainer_config)154 155 for uploaded_file in uploaded_files:156 image = Image.open(uploaded_file).convert("RGB")157 pred_class, probs = predict_image(model, transform, idx_to_class, image)158 159 st.subheader(f"Prediction: **{pred_class}**")160 st.bar_chart(pd.Series(probs, index=list(idx_to_class.values())))161 162 # GradCAM overlay using ExplanationGenerator163 img_tensor = transform(image).unsqueeze(0)164 heatmap, _ = explainer.gradcam(model, img_tensor, CONFIG["gradcam_target_layer"])165 plt_obj = explainer.create_side_by_side(image, heatmap, pred_class, "Unknown")166 167 # Display side-by-side result168 col1, col2 = st.columns(2)169 with col1:170 st.image(image, caption="Original Image", use_column_width=True)171 with col2:172 import io173 buf = io.BytesIO()174 plt_obj.savefig(buf, format="png", bbox_inches="tight")175 st.image(buf, caption="GradCAM Overlay", use_column_width=True)176 plt_obj.close()177 178# -------------------179# TAB 3: BATCH INFERENCE180# -------------------181with tabs[2]:182 st.header("Batch Prediction from CSV + Images ZIP")183 184 uploaded_csv = st.file_uploader("Upload CSV with filepaths", type=["csv"])185 uploaded_zip = st.file_uploader("Upload ZIP containing images", type=["zip"])186 187 if uploaded_csv and uploaded_zip:188 189 progress_bar = st.progress(0)190 status_text = st.empty()191 192 try:193 status_text.text("📂 Preparing files...")194 195 # Create a temp dir196 temp_dir = tempfile.mkdtemp()197 198 # Save and extract ZIP199 zip_path = os.path.join(temp_dir, "images.zip")200 with open(zip_path, "wb") as f:201 f.write(uploaded_zip.read())202 with zipfile.ZipFile(zip_path, 'r') as zip_ref:203 zip_ref.extractall(temp_dir)204 205 progress_bar.progress(20)206 status_text.text("✅ Images extracted")207 208 # Build map of filename -> actual extracted path209 file_map = {}210 for root, dirs, files in os.walk(temp_dir):211 for f in files:212 file_map[f] = os.path.join(root, f)213 214 # Read CSV215 df = pd.read_csv(uploaded_csv)216 progress_bar.progress(40)217 status_text.text("📄 CSV loaded")218 219 # Fix CSV paths to point to extracted folder220 def fix_path(p):221 filename = os.path.basename(p)222 if filename in file_map:223 return file_map[filename]224 else:225 raise FileNotFoundError(f"{filename} not found in uploaded ZIP.")226 227 df["filepath"] = df["filepath"].apply(fix_path)228 229 # Save updated CSV in temp dir230 temp_csv_path = os.path.join(temp_dir, "input.csv")231 df.to_csv(temp_csv_path, index=False)232 233 progress_bar.progress(60)234 status_text.text("⚙️ Configuring batch inference...")235 236 # Config object for BatchInference237 config_obj = type("Config", (object,), {238 "root_dir": temp_dir,239 "model_path": CONFIG["model_path"],240 "label_mapping_path": CONFIG["label_mapping_path"],241 "mean_std_path": CONFIG["mean_std_path"],242 "test_csv": temp_csv_path,243 "batch_size": 16,244 "num_workers": 0245 })()246 247 # Run inference248 batch_inf = BatchInference(config_obj)249 batch_inf.run()250 251 progress_bar.progress(90)252 253 # Show predictions254 pred_path = os.path.join(temp_dir, "batch_predictions.csv")255 if os.path.exists(pred_path):256 progress_bar.progress(100)257 status_text.text("✅ Inference complete!")258 st.success("✅ Predictions complete!")259 st.dataframe(pd.read_csv(pred_path).head())260 261 with open(pred_path, "rb") as f:262 st.download_button("Download CSV", f, file_name="batch_predictions.csv")263 264 except Exception as e:265 logger.exception(e)266 st.error(f"Error during batch inference: {str(e)}")267 status_text.text("❌ Error occurred. Please check logs.")268 progress_bar.progress(0)269 270# -------------------271# TAB 4: FINE-TUNING272# -------------------273with tabs[3]:274 st.header("🛠 Fine-tuning the Trained Model")275 st.markdown(276 "Upload **ImageNet style dataset** (train/ and val/ folders with species subfolders) "277 "to fine-tune the existing ResNet50 model."278 )279 280 uploaded_zip = st.file_uploader(281 "Upload dataset ZIP file (train/ and val/ inside)",282 type=["zip"]283 )284 285 # ---- User-selectable hyperparameters ----286 st.subheader("⚙️ Fine-tuning Settings")287 batch_size = st.number_input("Batch size", min_value=1, max_value=128, value=32, step=1)288 epochs = st.number_input("Epochs", min_value=1, max_value=100, value=5, step=1)289 unfreeze_backbone = st.checkbox("Unfreeze Backbone", value=False, help="Unfreeze all layers for deeper fine-tuning.")290 patience = st.number_input("Early Stopping Patience", min_value=1, max_value=20, value=3, step=1)291 learning_rate = st.number_input("Learning rate", min_value=1e-6, max_value=1.0, value=0.001, step=0.0001, format="%.4f")292 scheduler_patience = st.number_input("Scheduler Patience", min_value=1, max_value=10, value=2, step=1)293 scheduler_factor = st.number_input("Scheduler Factor", min_value=0.1, max_value=1.0, value=0.2, step=0.01)294 crop_size = st.number_input("Crop size", min_value=64, max_value=512, value=224, step=1)295 flip = st.checkbox("Random Horizontal Flip", value=True)296 brightness = st.slider("Brightness", 0.0, 1.0, 0.2, 0.05)297 contrast = st.slider("Contrast", 0.0, 1.0, 0.2, 0.05)298 saturation = st.slider("Saturation", 0.0, 1.0, 0.2, 0.05)299 hue = st.slider("Hue", 0.0, 0.5, 0.1, 0.01)300 301 if uploaded_zip is not None:302 tmp_dir = tempfile.mkdtemp()303 tmp_zip_path = os.path.join(tmp_dir, "dataset.zip")304 with open(tmp_zip_path, "wb") as f:305 f.write(uploaded_zip.getbuffer())306 307 with zipfile.ZipFile(tmp_zip_path, "r") as zip_ref:308 zip_ref.extractall(tmp_dir)309 310 # Paths311 train_dir = os.path.join(tmp_dir, "train")312 val_dir = os.path.join(tmp_dir, "val")313 314 if not (os.path.exists(train_dir) and os.path.exists(val_dir)):315 st.error("❌ Uploaded ZIP must contain 'train/' and 'val/' directories.")316 else:317 st.success("✅ Dataset extracted successfully.")318 319 # Config loading (modular)320 config_manager = ConfigurationManager()321 finetune_config = config_manager.get_fine_tuning_config(322 train_dir=train_dir,323 val_dir=val_dir,324 batch_size=batch_size,325 unfreeze_backbone=unfreeze_backbone, 326 epochs=epochs,327 patience=patience,328 learning_rate=learning_rate,329 scheduler_patience=scheduler_patience,330 scheduler_factor=scheduler_factor,331 crop_size=crop_size,332 flip=flip,333 brightness=brightness,334 contrast=contrast,335 saturation=saturation,336 hue=hue337 )338 339 if st.button("🚀 Start Fine-tuning"):340 try:341 # device info342 device_info = "cuda" if torch.cuda.is_available() else "cpu"343 st.info(f"💻 Training on: **{device_info.upper()}**")344 345 # Progress bar + real-time log placeholders346 progress_bar = st.progress(0)347 log_placeholder = st.empty()348 log_placeholder.text("🔄 Starting fine-tuning...")349 350 finetuner = FineTuner(finetune_config)351 352 # Hook for live logging353 def training_callback(epoch, total_epochs, train_loss, val_loss, val_acc):354 progress = int(((epoch + 1) / total_epochs) * 100)355 progress_bar.progress(progress)356 log_placeholder.markdown(357 f"**Epoch {epoch+1}/{total_epochs}** "358 f"Train Loss: `{train_loss:.4f}` "359 f"Val Loss: `{val_loss:.4f}` "360 f"Val Acc: `{val_acc*100:.2f}%`"361 )362 363 best_acc, train_losses, val_losses, val_accuracies = finetuner.run(callback=training_callback)364 if best_acc > 0:365 st.success(f"Fine-tuning complete! Best Validation Accuracy: {best_acc*100:.2f}% ")366 model_path = finetune_config.output_model_path367 label_map_path = finetune_config.output_label_mapping_path368 369 col1, col2 = st.columns(2)370 371 # Loss curves372 with col1:373 st.subheader("📉 Loss Curves")374 loss_df = pd.DataFrame({375 "Epoch": [str(e) for e in range(1, len(train_losses) + 1)], 376 "Train Loss": train_losses,377 "Validation Loss": val_losses378 })379 st.line_chart(loss_df, x="Epoch", y=["Train Loss", "Validation Loss"])380 381 # Accuracy curves382 with col2:383 st.subheader("✅ Validation Accuracy")384 acc_df = pd.DataFrame({385 "Epoch": [str(e) for e in range(1, len(val_accuracies) + 1)], 386 "Validation Accuracy (%)": [v * 100 for v in val_accuracies]387 })388 st.line_chart(acc_df, x="Epoch", y="Validation Accuracy (%)")389 with open(model_path, "rb") as f:390 st.download_button("⬇️ Download Fine-tuned Model (.pth)", f, file_name="fine_tuned_model.pth")391 392 with open(label_map_path, "rb") as f:393 st.download_button("⬇️ Download Label Mapping (.json)", f, file_name="label_mapping.json")394 395 else:396 st.error("⚠️ Fine-tuning stopped due to label mismatch.")397 except Exception as e:398 logger.exception(e)399 st.error(f"Error during fine-tuning: {str(e)}")400