| import os |
| import torch |
| import torch.nn as nn |
| import segmentation_models_pytorch as smp |
| import numpy as np |
| import albumentations as A |
| from albumentations.pytorch import ToTensorV2 |
| from app.config import config |
|
|
| |
| os.environ['TORCH_HOME'] = '/tmp/.cache/torch' |
| os.environ['HF_HOME'] = '/tmp/.cache/huggingface' |
| os.environ['TRANSFORMERS_CACHE'] = '/tmp/.cache/huggingface/transformers' |
|
|
| |
| os.makedirs('/tmp/.cache/torch', exist_ok=True) |
| os.makedirs('/tmp/.cache/huggingface', exist_ok=True) |
|
|
| def build_model(): |
| """Build and return the U-Net model""" |
| return smp.Unet( |
| encoder_name='efficientnet-b4', |
| encoder_weights='imagenet', |
| in_channels=3, |
| classes=1, |
| activation=None, |
| ) |
|
|
| def get_transforms(): |
| """Get image preprocessing transforms""" |
| return A.Compose([ |
| A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
| ToTensorV2(), |
| ]) |
|
|
| class ModelManager: |
| def __init__(self): |
| self.model = None |
| self.transforms = None |
| self.is_loaded = False |
|
|
| def load_model(self): |
| """Load the model and transforms""" |
| try: |
| self.model = build_model().to(config.DEVICE) |
| state_dict = torch.load(config.MODEL_PATH, map_location=config.DEVICE) |
| self.model.load_state_dict(state_dict) |
| self.model.eval() |
| self.transforms = get_transforms() |
| self.is_loaded = True |
| print(f"Model loaded successfully from {config.MODEL_PATH}") |
| return True |
| except Exception as e: |
| print(f"Model loading failed: {e}") |
| self.is_loaded = False |
| return False |
|
|
| def predict(self, image_np: np.ndarray): |
| """Make prediction on preprocessed image""" |
| if not self.is_loaded: |
| raise RuntimeError("Model not loaded") |
|
|
| |
| dummy_mask = np.zeros((config.IMG_HEIGHT, config.IMG_WIDTH), dtype=np.float32) |
| transformed = self.transforms(image=image_np, mask=dummy_mask) |
| input_tensor = transformed['image'].unsqueeze(0).to(config.DEVICE) |
|
|
| |
| with torch.no_grad(): |
| prediction_logits = self.model(input_tensor)[0, 0].cpu().numpy() |
|
|
| return prediction_logits |
|
|
| def postprocess_prediction(prediction: np.ndarray, threshold: float = 0.5) -> tuple: |
| """Convert model output to visualization images""" |
| |
| prob_map = 1 / (1 + np.exp(-prediction)) |
|
|
| |
| binary_mask = (prob_map > threshold).astype(np.uint8) * 255 |
|
|
| |
| confidence_score = float(np.max(prob_map)) |
| predicted_area_percentage = float(np.sum(prob_map > threshold) / (config.IMG_HEIGHT * config.IMG_WIDTH) * 100) |
|
|
| return prob_map, binary_mask, confidence_score, predicted_area_percentage |
|
|
| |
| model_manager = ModelManager() |
