Spaces:
Paused
Paused
| """ | |
| OralGuard AI - Prediction Utilities for Hugging Face Spaces | |
| Handles image preprocessing, prediction, and ensemble logic | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| from torchvision import transforms | |
| from PIL import Image | |
| import numpy as np | |
| import time | |
| import streamlit as st | |
| class OralGuardPredictor: | |
| """ | |
| Handles predictions for OralGuard AI models with on-demand loading | |
| """ | |
| def __init__(self, model_loader, img_size=224): | |
| self.model_loader = model_loader | |
| self.img_size = img_size | |
| self.device = model_loader.device | |
| self.class_names = ['Normal', 'OSCC'] | |
| # Define transforms | |
| self.transform = transforms.Compose([ | |
| transforms.Resize((img_size, img_size)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225] | |
| ) | |
| ]) | |
| def preprocess_image(self, image): | |
| """ | |
| Preprocess a PIL Image for model input | |
| Args: | |
| image: PIL Image | |
| Returns: | |
| Preprocessed tensor | |
| """ | |
| # Convert to RGB if necessary | |
| if image.mode != 'RGB': | |
| image = image.convert('RGB') | |
| # Apply transforms | |
| img_tensor = self.transform(image) | |
| # Add batch dimension | |
| img_tensor = img_tensor.unsqueeze(0) | |
| return img_tensor.to(self.device) | |
| def predict_single_model(self, image, model_type, magnification): | |
| """ | |
| Get prediction from a single model (loads on-demand) | |
| Args: | |
| image: PIL Image | |
| model_type: 'efficientnet', 'vgg16', or 'vit' | |
| magnification: '100x' or '400x' | |
| Returns: | |
| Dictionary with prediction results | |
| """ | |
| start_time = time.time() | |
| # Load model (downloads if needed) | |
| model = self.model_loader.load_model(model_type, magnification) | |
| # Preprocess image | |
| img_tensor = self.preprocess_image(image) | |
| # Make prediction | |
| with torch.no_grad(): | |
| outputs = model(img_tensor) | |
| probabilities = F.softmax(outputs, dim=1) | |
| confidence, predicted = torch.max(probabilities, 1) | |
| inference_time = time.time() - start_time | |
| # Convert to numpy | |
| probs = probabilities.cpu().numpy()[0] | |
| pred_class = predicted.cpu().item() | |
| conf_score = confidence.cpu().item() | |
| return { | |
| 'prediction': self.class_names[pred_class], | |
| 'confidence': float(conf_score), | |
| 'probabilities': { | |
| 'Normal': float(probs[0]), | |
| 'OSCC': float(probs[1]) | |
| }, | |
| 'inference_time': inference_time, | |
| 'model': f"{model_type}_{magnification}" | |
| } | |
| def predict_ensemble(self, image, magnification, weights=None): | |
| """ | |
| Get ensemble prediction from all three models | |
| Args: | |
| image: PIL Image | |
| magnification: '100x' or '400x' | |
| weights: Optional weights for models [efficientnet, vgg16, vit] | |
| Returns: | |
| Dictionary with ensemble prediction results | |
| """ | |
| start_time = time.time() | |
| # Default equal weights | |
| if weights is None: | |
| weights = [1/3, 1/3, 1/3] | |
| # Get predictions from all models | |
| model_types = ['efficientnet', 'vgg16', 'vit'] | |
| all_probs = [] | |
| individual_results = {} | |
| progress_bar = st.progress(0) | |
| status_text = st.empty() | |
| for idx, model_type in enumerate(model_types): | |
| status_text.text(f"Loading and running {model_type}...") | |
| result = self.predict_single_model(image, model_type, magnification) | |
| all_probs.append([ | |
| result['probabilities']['Normal'], | |
| result['probabilities']['OSCC'] | |
| ]) | |
| individual_results[model_type] = result | |
| progress_bar.progress((idx + 1) / len(model_types)) | |
| progress_bar.empty() | |
| status_text.empty() | |
| # Weighted average of probabilities | |
| all_probs = np.array(all_probs) | |
| weights = np.array(weights).reshape(-1, 1) | |
| ensemble_probs = (all_probs * weights).sum(axis=0) | |
| # Get prediction | |
| pred_class = np.argmax(ensemble_probs) | |
| confidence = float(ensemble_probs[pred_class]) | |
| inference_time = time.time() - start_time | |
| return { | |
| 'prediction': self.class_names[pred_class], | |
| 'confidence': confidence, | |
| 'probabilities': { | |
| 'Normal': float(ensemble_probs[0]), | |
| 'OSCC': float(ensemble_probs[1]) | |
| }, | |
| 'inference_time': inference_time, | |
| 'model': f'ensemble_{magnification}', | |
| 'individual_results': individual_results, | |
| 'ensemble_weights': { | |
| 'efficientnet': float(weights[0][0]), | |
| 'vgg16': float(weights[1][0]), | |
| 'vit': float(weights[2][0]) | |
| } | |
| } | |
| def predict_batch(self, images, model_type, magnification): | |
| """ | |
| Predict on a batch of images | |
| Args: | |
| images: List of PIL Images | |
| model_type: 'efficientnet', 'vgg16', 'vit', or 'ensemble' | |
| magnification: '100x' or '400x' | |
| Returns: | |
| List of prediction results | |
| """ | |
| results = [] | |
| progress_bar = st.progress(0) | |
| status_text = st.empty() | |
| for idx, image in enumerate(images): | |
| status_text.text(f"Processing image {idx+1}/{len(images)}...") | |
| try: | |
| if model_type == 'ensemble': | |
| # Use silent version to avoid nested progress bars | |
| result = self.predict_ensemble_silent(image, magnification) | |
| else: | |
| result = self.predict_single_model(image, model_type, magnification) | |
| result['image_index'] = idx | |
| result['status'] = 'success' | |
| results.append(result) | |
| except Exception as e: | |
| error_msg = str(e) | |
| st.error(f"❌ Error processing image {idx+1}: {error_msg}") | |
| results.append({ | |
| 'image_index': idx, | |
| 'status': 'error', | |
| 'error': error_msg, | |
| 'prediction': 'Error', | |
| 'confidence': 0.0 | |
| }) | |
| progress_bar.progress((idx + 1) / len(images)) | |
| progress_bar.empty() | |
| status_text.empty() | |
| # Show summary | |
| successful = sum(1 for r in results if r.get('status') == 'success') | |
| if successful < len(results): | |
| st.warning(f"⚠️ {len(results) - successful} images failed to process") | |
| return results | |
| def predict_ensemble_silent(self, image, magnification, weights=None): | |
| """ | |
| Get ensemble prediction without UI elements (for batch processing) | |
| Args: | |
| image: PIL Image | |
| magnification: '100x' or '400x' | |
| weights: Optional weights for models [efficientnet, vgg16, vit] | |
| Returns: | |
| Dictionary with ensemble prediction results | |
| """ | |
| start_time = time.time() | |
| # Default equal weights | |
| if weights is None: | |
| weights = [1/3, 1/3, 1/3] | |
| # Get predictions from all models | |
| model_types = ['efficientnet', 'vgg16', 'vit'] | |
| all_probs = [] | |
| individual_results = {} | |
| for model_type in model_types: | |
| result = self.predict_single_model(image, model_type, magnification) | |
| all_probs.append([ | |
| result['probabilities']['Normal'], | |
| result['probabilities']['OSCC'] | |
| ]) | |
| individual_results[model_type] = result | |
| # Weighted average of probabilities | |
| all_probs = np.array(all_probs) | |
| weights = np.array(weights).reshape(-1, 1) | |
| ensemble_probs = (all_probs * weights).sum(axis=0) | |
| # Get prediction | |
| pred_class = np.argmax(ensemble_probs) | |
| confidence = float(ensemble_probs[pred_class]) | |
| inference_time = time.time() - start_time | |
| return { | |
| 'prediction': self.class_names[pred_class], | |
| 'confidence': confidence, | |
| 'probabilities': { | |
| 'Normal': float(ensemble_probs[0]), | |
| 'OSCC': float(ensemble_probs[1]) | |
| }, | |
| 'inference_time': inference_time, | |
| 'model': f'ensemble_{magnification}', | |
| 'individual_results': individual_results, | |
| 'ensemble_weights': { | |
| 'efficientnet': float(weights[0][0]), | |
| 'vgg16': float(weights[1][0]), | |
| 'vit': float(weights[2][0]) | |
| } | |
| } | |
| def get_risk_level(self, prediction, confidence): | |
| """ | |
| Determine risk level based on prediction and confidence | |
| Args: | |
| prediction: 'Normal' or 'OSCC' | |
| confidence: confidence score (0-1) | |
| Returns: | |
| Risk level and color | |
| """ | |
| if prediction == 'Normal': | |
| if confidence >= 0.9: | |
| return 'Low Risk', 'green' | |
| elif confidence >= 0.7: | |
| return 'Low-Medium Risk', 'lightgreen' | |
| else: | |
| return 'Uncertain', 'orange' | |
| else: # OSCC | |
| if confidence >= 0.9: | |
| return 'High Risk', 'red' | |
| elif confidence >= 0.7: | |
| return 'Medium-High Risk', 'orange' | |
| else: | |
| return 'Uncertain', 'orange' | |
| def format_result_for_display(self, result): | |
| """ | |
| Format prediction result for display | |
| Args: | |
| result: Prediction result dictionary | |
| Returns: | |
| Formatted dictionary | |
| """ | |
| risk_level, risk_color = self.get_risk_level( | |
| result['prediction'], | |
| result['confidence'] | |
| ) | |
| formatted = { | |
| 'Prediction': result['prediction'], | |
| 'Confidence': f"{result['confidence']*100:.2f}%", | |
| 'Risk Level': risk_level, | |
| 'Risk Color': risk_color, | |
| 'Probability (Normal)': f"{result['probabilities']['Normal']*100:.2f}%", | |
| 'Probability (OSCC)': f"{result['probabilities']['OSCC']*100:.2f}%", | |
| 'Inference Time': f"{result['inference_time']:.3f}s", | |
| 'Model Used': result['model'] | |
| } | |
| # Add individual results for ensemble | |
| if 'individual_results' in result: | |
| formatted['Individual Models'] = {} | |
| for model, res in result['individual_results'].items(): | |
| formatted['Individual Models'][model] = { | |
| 'prediction': res['prediction'], | |
| 'confidence': f"{res['confidence']*100:.2f}%" | |
| } | |
| return formatted |