Spaces:
Sleeping
Sleeping
File size: 4,652 Bytes
6d1b0d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | 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()
|