main.py

# main.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
import tensorflow as tf
import numpy as np
import base64
import logging
import os
import sys
import time
from datetime import datetime
from logging.handlers import RotatingFileHandler

from inference.forest import predict_forest, build_input_tensor
from schemas import PredictRequest, PredictResponse

# =============================================================================
# LOGGING CONFIGURATION
# =============================================================================

os.makedirs("logs", exist_ok=True)

logger = logging.getLogger("forest_segmentation")
logger.setLevel(logging.DEBUG)

console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.DEBUG)
console_handler.setFormatter(
    logging.Formatter(
        "%(asctime)s | %(levelname)-8s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )
)

file_handler = RotatingFileHandler(
    "logs/server.log", maxBytes=10_000_000, backupCount=5, encoding="utf-8"
)
file_handler.setFormatter(console_handler.formatter)

logger.addHandler(console_handler)
logger.addHandler(file_handler)

logger.info("=" * 80)
logger.info("FOREST SEGMENTATION SERVER STARTING")
logger.info("=" * 80)

# =============================================================================
# INVERSION DETECTION
# =============================================================================

def detect_inversion(image_stack, confidence_map, ndvi_threshold=0.3):
    """
    Detect if model output is inverted using NDVI correlation.
    image_stack: (H, W, 9)
    confidence_map: (H, W)
    """
    ndvi = image_stack[:, :, 6]  # NDVI channel

    vegetation_mask = ndvi > ndvi_threshold

    veg_conf = (
        confidence_map[vegetation_mask].mean()
        if vegetation_mask.any() else 0.5
    )

    non_veg_conf = (
        confidence_map[~vegetation_mask].mean()
        if (~vegetation_mask).any() else 0.5
    )

    is_inverted = non_veg_conf > veg_conf
    correlation = veg_conf - non_veg_conf

    return bool(is_inverted), float(correlation)

# =============================================================================
# FASTAPI APP
# =============================================================================

app = FastAPI(
    title="Forest Segmentation API",
    description="Landsat 8 Forest Segmentation",
    version="1.0.0"
)

IMG_SIZE = 256
LANDSAT_BANDS = [
    "Blue", "Green", "Red",
    "NIR", "SWIR1", "SWIR2",
    "NDVI", "NDWI", "NBR"
]

# =============================================================================
# MIDDLEWARE
# =============================================================================

@app.middleware("http")
async def log_requests(request: Request, call_next):
    start = time.time()
    response = await call_next(request)
    duration = time.time() - start
    logger.info(
        f"{request.method} {request.url.path} | "
        f"{response.status_code} | {duration:.3f}s"
    )
    return response

# =============================================================================
# HEALTH
# =============================================================================

@app.get("/health")
def health():
    return {
        "status": "healthy",
        "timestamp": datetime.utcnow().isoformat()
    }

# =============================================================================
# PREDICT ENDPOINT (FIXED - CONTINUOUS VALUES)
# =============================================================================

@app.post("/predict", response_model=PredictResponse)
def predict(payload: PredictRequest):

    try:
        logger.info("[PREDICT] Request received")

        if not payload.bands:
            raise ValueError("No bands provided")

        # ---------------------------------------------------------------------
        # Decode bands
        # ---------------------------------------------------------------------
        decoded_bands = {}

        for band, data in payload.bands.items():
            if isinstance(data, str):
                raw = base64.b64decode(data)
                arr = np.frombuffer(raw, dtype=np.float32)
                side = int(np.sqrt(arr.size))
                decoded_bands[band] = arr.reshape((side, side))
            else:
                decoded_bands[band] = np.array(data, dtype=np.float32)

        logger.info(f"[PREDICT] Decoded {len(decoded_bands)} bands")

        # ---------------------------------------------------------------------
        # Build input tensor
        # ---------------------------------------------------------------------
        input_tensor = build_input_tensor(decoded_bands)  # (1, H, W, 9)
        input_stack = input_tensor[0]                     # (H, W, 9)

        # ---------------------------------------------------------------------
        # Run model (raw confidence)
        # ---------------------------------------------------------------------
        model = tf.keras.models.load_model(
            "models/Forest_Segmentation_Best.keras",
            compile=False
        )

        confidence_map = model.predict(
            input_tensor, verbose=0
        )[0, :, :, 0]

        # Log raw model output stats
        logger.info(
            f"[MODEL OUTPUT] Raw confidence: min={confidence_map.min():.4f}, "
            f"max={confidence_map.max():.4f}, mean={confidence_map.mean():.4f}"
        )

        # ---------------------------------------------------------------------
        # Inversion detection & correction
        # ---------------------------------------------------------------------
        is_inverted, corr = detect_inversion(
            input_stack, confidence_map
        )

        if is_inverted:
            logger.warning(
                f"[INVERSION] Detected | NDVI correlation={corr:.4f} | FIX APPLIED"
            )
            corrected_conf = 1.0 - confidence_map
        else:
            logger.info(
                f"[INVERSION] Not detected | NDVI correlation={corr:.4f}"
            )
            corrected_conf = confidence_map

        # ---------------------------------------------------------------------
        # Create masks (CONTINUOUS values for density visualization)
        # ---------------------------------------------------------------------
        # Use continuous confidence scaled to 0-255 (NOT binary!)
        mask_255 = (corrected_conf * 255).astype(np.uint8)
        inverted_mask_255 = (255 - mask_255).astype(np.uint8)

        # Calculate stats using threshold for percentage
        forest_percentage = float((corrected_conf > 0.5).sum() / corrected_conf.size * 100)
        forest_confidence = float(corrected_conf.mean())

        # Log mask stats to verify continuous values
        logger.info(
            f"[MASK] Range: [{mask_255.min()}, {mask_255.max()}] | "
            f"Unique values: {len(np.unique(mask_255))}"
        )

        logger.info(
            f"[PREDICT] Forest={forest_percentage:.2f}% | "
            f"Confidence={forest_confidence:.4f}"
        )

        # ---------------------------------------------------------------------
        # Response
        # ---------------------------------------------------------------------
        return {
            "mask": mask_255.flatten().tolist(),
            "inverted_mask": inverted_mask_255.flatten().tolist(),
            "forest_percentage": forest_percentage,
            "forest_confidence": forest_confidence,
            "mean_prediction": forest_confidence,
            "classes": {"forest": 1, "non_forest": 0},
            "model_info": {
                "name": "Forest_Segmentation_Best",
                "bands": LANDSAT_BANDS
            },
            "debug": {
                "was_inverted": is_inverted,
                "inversion_correlation": corr,
                "mask_min": int(mask_255.min()),
                "mask_max": int(mask_255.max()),
                "unique_values": int(len(np.unique(mask_255)))
            }
        }

    except ValueError as e:
        logger.error(f"[PREDICT] Validation error: {e}")
        raise HTTPException(status_code=400, detail=str(e))

    except Exception as e:
        logger.exception("[PREDICT] Inference failed")
        raise HTTPException(status_code=500, detail=str(e))

# =============================================================================
# STARTUP / SHUTDOWN
# =============================================================================

@app.on_event("startup")
async def startup():
    logger.info("=" * 80)
    logger.info("SERVER READY")
    logger.info("=" * 80)

@app.on_event("shutdown")
async def shutdown():
    logger.info("=" * 80)
    logger.info("SERVER SHUTDOWN")
    logger.info("=" * 80)