Spaces:
Sleeping
Sleeping
| import os | |
| import tensorflow as tf | |
| import numpy as np | |
| from PIL import Image | |
| import io | |
| import sys | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) | |
| from model.unet import build_unet | |
| from model.loss_metrics import get_metrics, bce_dice_loss | |
| from backend.utils.model_fetcher import download_if_missing | |
| # Force download if missing before initialization | |
| download_if_missing() | |
| # Path to the saved model weights | |
| MODEL_PATH = os.path.join(os.path.dirname(__file__), "../model/saved_models/oil_spill_unet_best.keras") | |
| IMG_SIZE = (256, 256) | |
| class OilSpillModel: | |
| def __init__(self): | |
| self.model = None | |
| self.load_model() | |
| def load_model(self): | |
| print(f"Attempting to load model from: {os.path.abspath(MODEL_PATH)}") | |
| if os.path.exists(MODEL_PATH): | |
| try: | |
| # Provide custom objects if model was saved with custom metrics/loss | |
| custom_objects = { | |
| 'bce_dice_loss': bce_dice_loss, | |
| 'dice_coef': get_metrics()[0], | |
| 'iou_metric': get_metrics()[1] | |
| } | |
| self.model = tf.keras.models.load_model(MODEL_PATH, custom_objects=custom_objects) | |
| print("Model loaded successfully.") | |
| except Exception as e: | |
| print(f"Failed to load model from file: {e}") | |
| self._build_stub_model() | |
| else: | |
| print("Trained model weights not found. Building an untrained stub model for development.") | |
| self._build_stub_model() | |
| def _build_stub_model(self): | |
| """Used for development when trained weights aren't available yet.""" | |
| # Builds architecture with random weights | |
| self.model = build_unet(input_shape=(256, 256, 3)) | |
| def predict(self, image_bytes): | |
| """ | |
| Takes raw image bytes, preprocesses, predicts, and returns the binary mask array and confidence score. | |
| """ | |
| # Load image | |
| img = Image.open(io.BytesIO(image_bytes)).convert("RGB") | |
| # Resize to network input shape | |
| img = img.resize(IMG_SIZE) | |
| # Convert to numpy array and normalize to [0, 1] | |
| img_array = np.array(img, dtype=np.float32) / 255.0 | |
| # Expand dimension to create a batch size of 1 | |
| img_array = np.expand_dims(img_array, axis=0) | |
| # Inference | |
| pred_mask = self.model.predict(img_array)[0] # Shape is (256, 256, 1) | |
| # Calculate confidence metric: | |
| # We average the probability of pixels that the network thinks are part of the spill (>0.5 probability) | |
| oil_pixels = pred_mask[pred_mask > 0.5] | |
| confidence = float(np.mean(oil_pixels)) if len(oil_pixels) > 0 else 0.0 | |
| # Threshold to create binary mask (255 for oil, 0 for background) | |
| binary_mask = (pred_mask > 0.5).astype(np.uint8) * 255 | |
| return binary_mask, confidence | |
| # Singleton prediction engine | |
| prediction_engine = OilSpillModel() | |