Spaces:
Paused
Paused
| """ | |
| OralGuard AI - Model Utilities for Hugging Face Spaces | |
| Handles Google Drive download and on-demand model loading | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import timm | |
| from efficientnet_pytorch import EfficientNet | |
| from torchvision import models | |
| from pathlib import Path | |
| import gdown | |
| import os | |
| import streamlit as st | |
| class OralGuardEfficientNet(nn.Module): | |
| """EfficientNet-B5 based model for oral cancer classification""" | |
| def __init__(self, num_classes=2, pretrained=False): | |
| super(OralGuardEfficientNet, self).__init__() | |
| if pretrained: | |
| self.backbone = EfficientNet.from_pretrained('efficientnet-b5') | |
| else: | |
| self.backbone = EfficientNet.from_name('efficientnet-b5') | |
| num_features = self.backbone._fc.in_features | |
| self.backbone._fc = nn.Sequential( | |
| nn.Dropout(0.3), | |
| nn.Linear(num_features, 512), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(512, num_classes) | |
| ) | |
| def forward(self, x): | |
| return self.backbone(x) | |
| class OralGuardVGG16(nn.Module): | |
| """VGG16 with Batch Normalization for oral cancer classification""" | |
| def __init__(self, num_classes=2, pretrained=False): | |
| super(OralGuardVGG16, self).__init__() | |
| self.backbone = models.vgg16_bn(pretrained=pretrained) | |
| num_features = self.backbone.classifier[6].in_features | |
| self.backbone.classifier[6] = nn.Sequential( | |
| nn.Dropout(0.3), | |
| nn.Linear(num_features, 512), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(512, num_classes) | |
| ) | |
| def forward(self, x): | |
| return self.backbone(x) | |
| class OralGuardViT(nn.Module): | |
| """Vision Transformer (ViT-B/16) for oral cancer classification""" | |
| def __init__(self, num_classes=2, pretrained=False): | |
| super(OralGuardViT, self).__init__() | |
| self.backbone = timm.create_model( | |
| 'vit_base_patch16_224', | |
| pretrained=pretrained, | |
| num_classes=0 | |
| ) | |
| embed_dim = self.backbone.embed_dim | |
| self.classifier = nn.Sequential( | |
| nn.Linear(embed_dim, 512), | |
| nn.ReLU(), | |
| nn.Dropout(0.3), | |
| nn.Linear(512, 256), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(256, num_classes) | |
| ) | |
| def forward(self, x): | |
| features = self.backbone(x) | |
| return self.classifier(features) | |
| class GoogleDriveModelLoader: | |
| """ | |
| Handles Google Drive model downloading and loading with caching | |
| """ | |
| def __init__(self, cache_dir="./models_cache"): | |
| self.cache_dir = Path(cache_dir) | |
| self.cache_dir.mkdir(parents=True, exist_ok=True) | |
| self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| self.loaded_models = {} | |
| self.class_names = ['Normal', 'OSCC'] | |
| # Google Drive file IDs - will be set from Hugging Face secrets or config | |
| self.gdrive_file_ids = self._get_gdrive_ids() | |
| def _get_gdrive_ids(self): | |
| """ | |
| Get Google Drive file IDs from Hugging Face secrets or environment | |
| Format: https://drive.google.com/file/d/FILE_ID/view?usp=sharing | |
| """ | |
| # Try to get from Streamlit secrets first, fallback to env vars | |
| try: | |
| return { | |
| 'efficientnet_100x': st.secrets.get("EFFICIENTNET_100X_ID", os.getenv("EFFICIENTNET_100X_ID")), | |
| 'efficientnet_400x': st.secrets.get("EFFICIENTNET_400X_ID", os.getenv("EFFICIENTNET_400X_ID")), | |
| 'vgg16_100x': st.secrets.get("VGG16_100X_ID", os.getenv("VGG16_100X_ID")), | |
| 'vgg16_400x': st.secrets.get("VGG16_400X_ID", os.getenv("VGG16_400X_ID")), | |
| 'vit_100x': st.secrets.get("VIT_100X_ID", os.getenv("VIT_100X_ID")), | |
| 'vit_400x': st.secrets.get("VIT_400X_ID", os.getenv("VIT_400X_ID")), | |
| } | |
| except: | |
| # Return empty dict if secrets not configured | |
| return {} | |
| def download_model_from_gdrive(self, file_id, model_name, magnification): | |
| """ | |
| Download model from Google Drive if not already cached | |
| """ | |
| # Ensure cache directory exists | |
| self.cache_dir.mkdir(parents=True, exist_ok=True) | |
| # Create filename | |
| filename = f"{model_name}_{magnification}.pth" | |
| filepath = self.cache_dir / filename | |
| # Check if already downloaded | |
| if filepath.exists(): | |
| st.info(f"✅ Using cached model: {filename}") | |
| return filepath | |
| # Check file_id | |
| if not file_id: | |
| raise ValueError(f"Google Drive file ID not configured for {model_name}_{magnification}") | |
| # Show download info | |
| download_msg = st.empty() | |
| download_msg.info(f"⬇️ Downloading {filename} from Google Drive... This may take a few minutes.") | |
| try: | |
| # Construct download URL | |
| url = f"https://drive.google.com/uc?id={file_id}" | |
| # Create output path string | |
| output = str(filepath.absolute()) | |
| # Download with gdown (quiet mode to reduce jitter) | |
| import gdown | |
| #gdown.download(url, output, quiet=True) | |
| # Download with gdown WITH output to see what's happening | |
| st.write(f"📥 Downloading from: {url}") | |
| st.write(f"📂 Saving to: {output}") | |
| st.write(f"📁 Cache dir: {self.cache_dir.absolute()}") | |
| # Download with verbose output | |
| import sys | |
| from io import StringIO | |
| old_stdout = sys.stdout | |
| sys.stdout = buffer = StringIO() | |
| try: | |
| result = gdown.download(url, output, quiet=False) | |
| download_output = buffer.getvalue() | |
| st.code(download_output) | |
| st.write(f"✅ gdown returned: {result}") | |
| finally: | |
| sys.stdout = old_stdout | |
| # NOW check what actually happened | |
| st.write(f"📂 Checking: {output}") | |
| st.write(f" - Exists: {os.path.exists(output)}") | |
| if os.path.exists(output): | |
| st.write(f" - Size: {os.path.getsize(output)} bytes") | |
| # Also check the cache directory | |
| st.write(f"📁 Cache directory contents:") | |
| for f in self.cache_dir.glob("*"): | |
| st.write(f" - {f.name}: {f.stat().st_size} bytes") | |
| # Verify file was downloaded | |
| if not filepath.exists(): | |
| raise FileNotFoundError(f"Download completed but file not found at {filepath}") | |
| download_msg.success(f"✅ Downloaded {filename} successfully!") | |
| # Clear message after 2 seconds | |
| import time | |
| time.sleep(2) | |
| download_msg.empty() | |
| return filepath | |
| except Exception as e: | |
| download_msg.error(f"❌ Error downloading {filename}: {str(e)}") | |
| # Clean up partial download | |
| if filepath.exists(): | |
| filepath.unlink() | |
| raise | |
| def load_model(self, model_type, magnification): | |
| """ | |
| Load a specific model (downloads if necessary) | |
| """ | |
| model_key = f"{model_type}_{magnification}" | |
| # Return if already loaded | |
| if model_key in self.loaded_models: | |
| return self.loaded_models[model_key] | |
| # Get file ID | |
| file_id = self.gdrive_file_ids.get(model_key) | |
| # Download model if needed | |
| model_path = self.download_model_from_gdrive(file_id, model_type, magnification) | |
| # CREATE MODEL INSTANCE FIRST | |
| if model_type == 'efficientnet': | |
| model = OralGuardEfficientNet(num_classes=2, pretrained=False) | |
| elif model_type == 'vgg16': | |
| model = OralGuardVGG16(num_classes=2, pretrained=False) | |
| elif model_type == 'vit': | |
| model = OralGuardViT(num_classes=2, pretrained=False) | |
| else: | |
| raise ValueError(f"Unknown model type: {model_type}") | |
| # CRITICAL FIX: Convert Path to absolute string | |
| model_path_str = str(model_path.absolute()) | |
| # Load weights | |
| st.info(f"🔄 Loading {model_key} into memory...") | |
| try: | |
| checkpoint = torch.load(model_path_str, map_location=self.device) | |
| # Handle different checkpoint formats | |
| if isinstance(checkpoint, dict): | |
| if 'model_state_dict' in checkpoint: | |
| model.load_state_dict(checkpoint['model_state_dict']) | |
| elif 'state_dict' in checkpoint: | |
| model.load_state_dict(checkpoint['state_dict']) | |
| else: | |
| model.load_state_dict(checkpoint) | |
| else: | |
| model.load_state_dict(checkpoint) | |
| # Move to device and set to eval mode | |
| model = model.to(self.device) | |
| model.eval() | |
| # Cache the loaded model | |
| self.loaded_models[model_key] = model | |
| st.success(f"✅ {model_key} loaded successfully!") | |
| return model | |
| except Exception as e: | |
| st.error(f"❌ Failed to load model: {str(e)}") | |
| st.error(f"Model path: {model_path_str}") | |
| st.error(f"File exists: {os.path.exists(model_path_str)}") | |
| raise | |
| def get_model_info(self): | |
| """Get information about loaded models""" | |
| info = {} | |
| for key, model in self.loaded_models.items(): | |
| total_params = sum(p.numel() for p in model.parameters()) | |
| trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) | |
| info[key] = { | |
| 'total_parameters': f"{total_params:,}", | |
| 'trainable_parameters': f"{trainable_params:,}", | |
| 'device': str(next(model.parameters()).device), | |
| 'loaded': True | |
| } | |
| return info | |
| def clear_cache(self): | |
| """Clear loaded models from memory""" | |
| self.loaded_models.clear() | |
| torch.cuda.empty_cache() if torch.cuda.is_available() else None | |
| st.success("✅ Model cache cleared!") | |
| def get_model_display_name(model_type): | |
| """Get display name for model type""" | |
| names = { | |
| 'efficientnet': 'EfficientNet-B5', | |
| 'vgg16': 'VGG16-BN', | |
| 'vit': 'ViT-B/16', | |
| 'ensemble': 'Ensemble (All Models)' | |
| } | |
| return names.get(model_type, model_type) | |
| def get_magnification_display_name(mag): | |
| """Get display name for magnification""" | |
| names = { | |
| '100x': '100× Magnification', | |
| '400x': '400× Magnification' | |
| } | |
| return names.get(mag, mag) | |
| def get_model_size_info(): | |
| """Get expected model sizes""" | |
| return { | |
| 'efficientnet': '340 MB per model', | |
| 'vgg16': '1.6 GB per model', | |
| 'vit': '1.0 GB per model', | |
| 'total_per_magnification': '~2.94 GB', | |
| 'total_all': '~5.88 GB' | |
| } |