Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| import albumentations as A | |
| from albumentations.pytorch import ToTensorV2 | |
| from .architectures import AuxModel, BiomassModel | |
| from ..core.config import settings | |
| class ModelManager: | |
| """Singleton pattern or persistent instance to manage models and memory out of route logic""" | |
| def __init__(self): | |
| self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| self.aux_models = [] | |
| self.main_models = [] | |
| self.transform = A.Compose([ | |
| A.Resize(settings.img_size, settings.img_size), | |
| A.Normalize(), | |
| ToTensorV2() | |
| ]) | |
| self.is_loaded = False | |
| def load_models(self): | |
| if self.is_loaded: | |
| return | |
| print(f"Loading models to {self.device}...") | |
| # Load Aux Models | |
| for seed in settings.aux_seeds: | |
| for fold in range(5): | |
| path = os.path.join(settings.aux_model_dir, f'best_aux_only_seed{seed}_fold{fold}.pth') | |
| if not os.path.exists(path): | |
| path_alt = os.path.join(settings.aux_model_dir, f'best_aux_seed{seed}_fold{fold}.pth') | |
| if os.path.exists(path_alt): | |
| path = path_alt | |
| else: continue | |
| ckpt = torch.load(path, map_location=self.device, weights_only=False) | |
| model = AuxModel(settings.model_name).to(self.device).eval() | |
| model.load_state_dict(ckpt['model_state_dict']) | |
| self.aux_models.append((model, ckpt)) | |
| # Load Main Models | |
| for seed in settings.seeds: | |
| for fold in settings.fold_weights.keys(): | |
| path = os.path.join(settings.main_model_dir, f'best_model_seed{seed}_fold{fold}.pth') | |
| if not os.path.exists(path): continue | |
| ckpt = torch.load(path, map_location=self.device, weights_only=False) | |
| model = BiomassModel(settings.model_name, img_size=settings.img_size).to(self.device).eval() | |
| model.load_state_dict(ckpt['model_state_dict'] if 'model_state_dict' in ckpt else ckpt) | |
| self.main_models.append((model, ckpt, settings.fold_weights[fold])) | |
| self.is_loaded = True | |
| print(f"Loaded {len(self.aux_models)} Aux Models and {len(self.main_models)} Main Models successfully.") | |
| def predict(self, image_bytes: bytes) -> dict: | |
| if not self.is_loaded: | |
| raise RuntimeError("Models are not loaded.") | |
| # Decode Image | |
| nparr = np.frombuffer(image_bytes, np.uint8) | |
| img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
| if img is None: | |
| raise ValueError("Could not decode image.") | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| # Transform | |
| tensor_img = self.transform(image=img)['image'].unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| # ---- Stage 1: Aux Prediction ---- | |
| aux_preds_accum = 0 | |
| for model, ckpt in self.aux_models: | |
| res = model(tensor_img).cpu().numpy() | |
| if 'tab_scaler' in ckpt: | |
| res = ckpt['tab_scaler'].inverse_transform(res) | |
| aux_preds_accum += res | |
| # Average aux predictions | |
| predicted_tabular = aux_preds_accum / max(1, len(self.aux_models)) | |
| # ---- Stage 2: Main Prediction ---- | |
| final_biomass_accum = 0 | |
| total_w = 0 | |
| for model, ckpt, weight in self.main_models: | |
| tab_input = predicted_tabular.copy() | |
| if ckpt.get('tabular_scaler'): | |
| tab_input = ckpt['tabular_scaler'].transform(tab_input) | |
| tab_tensor = torch.tensor(tab_input, dtype=torch.float32).to(self.device) | |
| res = model(tensor_img, tab_tensor).cpu().numpy() | |
| if ckpt.get('target_scaler'): | |
| res = ckpt['target_scaler'].inverse_transform(res) | |
| final_biomass_accum += (res * weight) | |
| total_w += weight | |
| if total_w == 0: | |
| raise RuntimeError("No Main Models available for prediction.") | |
| final_preds = final_biomass_accum / total_w | |
| # Format results (equivalent to np.maximum(0, val)) | |
| results = {} | |
| for j, target in enumerate(settings.targets): | |
| results[target] = max(0.0, float(final_preds[0, j])) | |
| return results | |
| model_manager = ModelManager() | |