Spaces:
Runtime error
Runtime error
File size: 1,735 Bytes
1c1956b | 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 | """
Unified Inference Engine for MedAdapt-SAM Dashboard
"""
import torch
import torch.nn as nn
import numpy as np
from typing import Dict, Any
class MedAdaptEngine:
"""Unified engine to handle different model types for the dashboard"""
def __init__(self, checkpoints_path: str = "checkpoints/"):
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.checkpoints_path = checkpoints_path
self.models = {}
def load_unet(self):
from models.unet import UNet
model = UNet(in_channels=3, out_channels=3)
# Load weights if exist
return model.to(self.device).eval()
def load_sam_adapter(self):
from models.sam_adapter import SAMAdapter
model = SAMAdapter(model_type='vit_b')
# Load weights if exist
return model.to(self.device).eval()
def predict(self, image_np: np.ndarray, model_type: str) -> Dict[str, np.ndarray]:
"""Generic prediction wrapper"""
# 1. Preprocess
img_tensor = torch.from_numpy(image_np).permute(2, 0, 1).float().divide(255.0).unsqueeze(0).to(self.device)
# 2. Logic based on model_type
# (This is where the actual multi-phase model logic connects)
# Placeholder for real output
# In production: with torch.no_grad(): out = self.models[model_type](img_tensor)
# Returning dummy masks for structural validation
h, w = image_np.shape[:2]
return {
'ET': (np.random.rand(h, w) > 0.95).astype(np.float32),
'TC': (np.random.rand(h, w) > 0.90).astype(np.float32),
'WT': (np.random.rand(h, w) > 0.85).astype(np.float32)
}
|