ndvi_pred only

app.py

@@ -0,0 +1,41 @@
+# app.py
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import StreamingResponse, JSONResponse
+from ndvi_predictor import load_model, normalize_rgb, predict_ndvi, create_visualization
+from PIL import Image
+import numpy as np
+import io
+import os
+
+app = FastAPI()
+model = load_model("ndvi_model.keras")
+
+@app.post("/predict/")
+async def predict_ndvi_api(file: UploadFile = File(...)):
+    img = Image.open(await file.read()).convert("RGB")
+    rgb_np = np.array(img)
+    rgb_norm = normalize_rgb(rgb_np)
+    ndvi = predict_ndvi(model, rgb_norm)
+
+    # Generate visual output
+    vis_image = create_visualization(rgb_np, ndvi)
+    vis_filename = "ndvi_visualization.png"
+
+    # Generate GeoTIFF in memory
+    geotiff_buffer = io.BytesIO()
+    import rasterio
+    profile = {
+        'driver': 'GTiff',
+        'height': ndvi.shape[0],
+        'width': ndvi.shape[1],
+        'count': 1,
+        'dtype': 'float32'
+    }
+    with rasterio.open(geotiff_buffer, 'w', **profile) as dst:
+        dst.write(ndvi, 1)
+    geotiff_buffer.seek(0)
+
+    return {
+        "ndvi_tiff": StreamingResponse(geotiff_buffer, media_type="image/tiff", headers={"Content-Disposition": "attachment; filename=predicted_ndvi.tif"}),
+        "visualization": StreamingResponse(vis_image, media_type="image/png", headers={"Content-Disposition": "attachment; filename=ndvi_visualization.png"})
+    }

dockerfile

@@ -0,0 +1,25 @@
+FROM python:3.10
+
+# Avoid interactive prompts
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    gdal-bin \
+    python3-gdal \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set work directory
+WORKDIR /code
+
+# Install Python packages
+COPY requirements.txt .
+RUN pip install --upgrade pip && pip install -r requirements.txt
+
+# Copy all files
+COPY . .
+
+# Launch FastAPI app with Uvicorn
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

ndvi_predictor.py

@@ -0,0 +1,76 @@
+# ndvi_predictor.py
+import os
+os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+os.environ["SM_FRAMEWORK"] = "tf.keras"
+
+import segmentation_models as sm
+import tensorflow as tf
+import numpy as np
+import rasterio
+import matplotlib.pyplot as plt
+from PIL import Image
+import io
+
+def load_model(model_path):
+    return tf.keras.models.load_model(model_path, compile=False)
+
+def normalize_rgb(rgb):
+    rgb_norm = rgb.copy().astype(np.float32)
+    for b in range(3):
+        band = rgb_norm[:, :, b]
+        min_val, max_val = np.percentile(band, [1, 99])
+        if min_val < max_val:
+            rgb_norm[:, :, b] = np.clip((band - min_val) / (max_val - min_val), 0, 1)
+    return rgb_norm
+
+def predict_ndvi(model, rgb_np):
+    height, width = rgb_np.shape[:2]
+    tile_size = 512
+    stride = int(tile_size * 0.7)
+
+    ndvi_pred = np.zeros((height, width), dtype=np.float32)
+    weight_map = np.zeros((height, width), dtype=np.float32)
+
+    if height < tile_size or width < tile_size:
+        pad_height = max(0, tile_size - height)
+        pad_width = max(0, tile_size - width)
+        rgb_padded = np.pad(rgb_np, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')
+        height_padded, width_padded = rgb_padded.shape[0], rgb_padded.shape[1]
+    else:
+        rgb_padded = rgb_np
+        height_padded, width_padded = height, width
+
+    for i in range(0, height_padded - tile_size + 1, stride):
+        for j in range(0, width_padded - tile_size + 1, stride):
+            tile = rgb_padded[i:i+tile_size, j:j+tile_size, :]
+            y, x = np.mgrid[0:tile_size, 0:tile_size]
+            weights = np.minimum(np.minimum(x, tile_size - x - 1), np.minimum(y, tile_size - y - 1))
+            weights = np.clip(weights, 0, 50) / 50
+            tile_pred = model.predict(np.expand_dims(tile, axis=0), verbose=0)[0, :, :, 0]
+            valid_height = min(tile_size, height - i)
+            valid_width = min(tile_size, width - j)
+            ndvi_pred[i:i+valid_height, j:j+valid_width] += tile_pred[:valid_height, :valid_width] * weights[:valid_height, :valid_width]
+            weight_map[i:i+valid_height, j:j+valid_width] += weights[:valid_height, :valid_width]
+
+    mask = weight_map > 0
+    ndvi_pred[mask] = ndvi_pred[mask] / weight_map[mask]
+    return ndvi_pred
+
+def create_visualization(rgb, ndvi):
+    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
+    rgb_disp = np.clip(rgb / 255 if rgb.max() > 1 else rgb, 0, 1)
+    axes[0].imshow(rgb_disp)
+    axes[0].set_title("RGB Input")
+    axes[0].axis("off")
+
+    im = axes[1].imshow(ndvi, cmap='RdYlGn', vmin=-1, vmax=1)
+    axes[1].set_title("Predicted NDVI")
+    axes[1].axis("off")
+    fig.colorbar(im, ax=axes[1])
+
+    buf = io.BytesIO()
+    plt.tight_layout()
+    plt.savefig(buf, format="png")
+    plt.close(fig)
+    buf.seek(0)
+    return buf

requirements.txt

@@ -0,0 +1,9 @@
+tensorflow
+segmentation-models
+rasterio
+numpy
+matplotlib
+Pillow
+tqdm
+fastapi
+uvicorn