FastAPI for Landsat vessel detection

Dockerfile

@@ -1,12 +1,8 @@
 FROM pytorch/pytorch:2.4.0-cuda11.8-cudnn9-runtime@sha256:58a28ab734f23561aa146fbaf777fb319a953ca1e188832863ed57d510c9f197
 
-# TEMPORARY Until RSLEARN Is Public
-ARG GIT_USERNAME
-ARG GIT_TOKEN
-
 RUN apt update
 RUN apt install -y libpq-dev ffmpeg libsm6 libxext6 git
-RUN git clone https://${GIT_USERNAME}:${GIT_TOKEN}@github.com/allenai/rslearn.git /opt/rslearn_projects/rslearn
+RUN git clone https://github.com/allenai/rslearn.git /opt/rslearn_projects/rslearn
 RUN pip install -r /opt/rslearn_projects/rslearn/requirements.txt
 RUN pip install -r /opt/rslearn_projects/rslearn/extra_requirements.txt
 COPY requirements.txt /opt/rslearn_projects/requirements.txt

landsat/recheck_landsat_labels/phase123_config.yaml

@@ -52,6 +52,8 @@ data:
               allow_invalid: true
               skip_unknown_categories: true
               prob_property: "prob"
+              positive_class: "correct"
+              positive_class_threshold: 0.85
         input_mapping:
           class:
             label: "targets"

one_off_projects/minderoo/classify/add_train_fp_examples.py

@@ -28,8 +28,12 @@
 from rslearn.utils.vector_format import GeojsonVectorFormat
 
 if __name__ == "__main__":
-    json_fname = "/multisat/datasets/dvim/train_outputs/2024-07-30_0.5/20240811_annotations.json"
-    big_crop_dir = "/multisat/datasets/dvim/train_outputs/2024-07-30_0.5/annotate_crops/"
+    json_fname = (
+        "/multisat/datasets/dvim/train_outputs/2024-07-30_0.5/20240811_annotations.json"
+    )
+    big_crop_dir = (
+        "/multisat/datasets/dvim/train_outputs/2024-07-30_0.5/annotate_crops/"
+    )
     out_dir = "/multisat/datasets/dvim/rslearn_classify/"
 
     # We are just using fake projection for training this model.

requirements.txt

@@ -1,3 +1,8 @@
 beaker-py
 python-dotenv
 pytest
+uvicorn
+fastapi
+pydantic
+typing-extensions
+ruff

rslp/landsat_vessels/Dockerfile

@@ -0,0 +1,12 @@
+# Base image
+FROM base-image:latest
+
+# Environment variables
+ENV PYTHONPATH="/opt/rslearn_projects:${PYTHONPATH}"
+ENV LANDSAT_PORT=5555
+
+# Make port 5555 available to the world outside this container
+EXPOSE $LANDSAT_PORT
+
+# Run app.py when the container launches
+CMD ["python3", "rslp/landsat_vessels/api_main.py"]

rslp/landsat_vessels/api_main.py

@@ -0,0 +1,99 @@
+"""Landsat Vessel Detection Service."""
+
+from __future__ import annotations
+
+import logging
+import multiprocessing
+import os
+
+import uvicorn
+from fastapi import FastAPI, Response
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+from rslp.landsat_vessels import predict_pipeline
+
+app = FastAPI()
+
+# Set up the logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+LANDSAT_HOST = "0.0.0.0"
+LANDSAT_PORT = 5555
+
+
+class FormattedPrediction(TypedDict):
+    """Formatted prediction for a single vessel detection."""
+
+    latitude: float
+    longitude: float
+    score: float
+    rgb_fname: str
+    b8_fname: str
+
+
+class LandsatResponse(BaseModel):
+    """Response object for vessel detections."""
+
+    status: list[str]
+    predictions: list[FormattedPrediction]
+
+
+class LandsatRequest(BaseModel):
+    """Request object for vessel detections."""
+
+    scene_id: str | None = None
+    image_files: dict[str, str] | None = None
+    crop_path: str
+    scratch_path: str
+    json_path: str
+
+
+@app.on_event("startup")
+async def rslp_init() -> None:
+    """Landsat Vessel Service Initialization."""
+    logger.info("Initializing")
+    multiprocessing.set_start_method("forkserver", force=True)
+    multiprocessing.set_forkserver_preload(
+        [
+            "rslp.utils.rslearn.materialize_dataset",
+            "rslp.utils.rslearn.run_model_predict",
+        ]
+    )
+
+
+@app.get("/")
+async def home() -> dict:
+    """Returns a simple message to indicate the service is running."""
+    return {"message": "Landsat Detections App"}
+
+
+@app.post("/detections", response_model=LandsatResponse)
+async def get_detections(info: LandsatRequest, response: Response) -> LandsatResponse:
+    """Returns vessel detections Response object for a given Request object."""
+    try:
+        logger.info(f"Received request with scene_id: {info.scene_id}")
+        json_data = predict_pipeline(
+            crop_path=info.crop_path,
+            scene_id=info.scene_id,
+            image_files=info.image_files,
+            scratch_path=info.scratch_path,
+            json_path=info.json_path,
+        )
+        return LandsatResponse(status=["success"], predictions=json_data)
+    except ValueError as e:
+        logger.error(f"Value error during prediction pipeline: {e}")
+        return LandsatResponse(status=["error"], predictions=[])
+    except Exception as e:
+        logger.error(f"Unexpected error during prediction pipeline: {e}")
+        return LandsatResponse(status=["error"], predictions=[])
+
+
+if __name__ == "__main__":
+    uvicorn.run(
+        "api_main:app",
+        host=os.getenv("LANDSAT_HOST", default="0.0.0.0"),
+        port=int(os.getenv("LANDSAT_PORT", default=5555)),
+        proxy_headers=True,
+    )

rslp/landsat_vessels/docker-compose.yaml

@@ -0,0 +1,33 @@
+version: "3.9"
+
+services:
+  # Define the base image
+  base-image:
+    build:
+      context: ../..
+      dockerfile: Dockerfile
+    image: base-image:latest   # Tag it as "base-image"
+
+  # Define the landsat-vessels service
+  landsat-vessels:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    shm_size: '10G'  # This adds the shared memory size
+    depends_on:
+      - base-image
+    ports:
+      - "5555:5555"
+    environment:
+      - RSLP_BUCKET
+      - S3_ACCESS_KEY_ID
+      - S3_SECRET_ACCESS_KEY
+      - AWS_ACCESS_KEY_ID
+      - AWS_SECRET_ACCESS_KEY
+      - NVIDIA_VISIBLE_DEVICES=all  # Make all GPUs visible
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]  # Ensure this service can access GPUs
+    runtime: nvidia  # Use the NVIDIA runtime

rslp/landsat_vessels/predict_pipeline.py

@@ -1,7 +1,9 @@
 """Landsat vessel prediction pipeline."""
 
 import json
+import time
 from datetime import datetime, timedelta
+from typing import Any
 
 import numpy as np
 import rasterio
@@ -181,12 +183,12 @@ def run_classifier(
 
 
 def predict_pipeline(
-    scratch_path: str,
-    json_path: str,
     crop_path: str,
+    scratch_path: str | None = None,
+    json_path: str | None = None,
     image_files: dict[str, str] | None = None,
     scene_id: str | None = None,
-):
+) -> dict[str, Any]:
     """Run the Landsat vessel prediction pipeline.
 
     This inputs a Landsat scene (consisting of per-band GeoTIFFs) and produces the
@@ -202,6 +204,9 @@ def predict_pipeline(
         scene_id: Landsat scene ID. Exactly one of image_files or scene_id should be
             specified.
     """
+    start_time = time.time()  # Start the timer
+    time_profile = {}
+
     ds_path = UPath(scratch_path)
     ds_path.mkdir(parents=True, exist_ok=True)
 
@@ -260,15 +265,25 @@ def predict_pipeline(
             dst_geom.time_range[1] + timedelta(minutes=30),
         )
 
+    time_profile["setup"] = time.time() - start_time
+
     # Run pipeline.
+    step_start_time = time.time()
+    print("run detector")
     detections = get_vessel_detections(
         ds_path, projection, scene_bounds, time_range=time_range
     )
+    time_profile["get_vessel_detections"] = time.time() - step_start_time
+
+    step_start_time = time.time()
+    print("run classifier")
     detections = run_classifier(ds_path, detections, time_range=time_range)
+    time_profile["run_classifier"] = time.time() - step_start_time
 
     # Write JSON and crops.
-    json_path = UPath(json_path)
+    step_start_time = time.time()
     crop_path = UPath(crop_path)
+    crop_path.mkdir(parents=True, exist_ok=True)
 
     json_data = []
     for idx, detection in enumerate(detections):
@@ -326,5 +341,18 @@ def predict_pipeline(
             )
         )
 
-    with json_path.open("w") as f:
-        json.dump(json_data, f)
+    time_profile["write_json_and_crops"] = time.time() - step_start_time
+
+    elapsed_time = time.time() - start_time  # Calculate elapsed time
+    time_profile["total"] = elapsed_time
+
+    if json_path:
+        json_path = UPath(json_path)
+        with json_path.open("w") as f:
+            json.dump(json_data, f)
+
+    print(f"Prediction pipeline completed in {elapsed_time:.2f} seconds")
+    for step, duration in time_profile.items():
+        print(f"{step} took {duration:.2f} seconds")
+
+    return json_data

rslp/utils/rslearn.py

@@ -34,12 +34,14 @@ def materialize_dataset(ds_path: UPath, group: str | None = None, workers: int =
         dataset,
         workers=workers,
         group=group,
+        use_initial_job=False,
     )
     apply_on_windows(
         MaterializeHandler(),
         dataset,
         workers=workers,
         group=group,
+        use_initial_job=False,
     )
 
 

tests/integration/landsat_vessels/test_fastapi.py

@@ -0,0 +1,26 @@
+import tempfile
+
+from fastapi.testclient import TestClient
+
+from rslp.landsat_vessels.api_main import app
+
+client = TestClient(app)
+
+
+def test_singapore_dense_scene():
+    # LC08_L1TP_125059_20240913_20240920_02_T1 is a scene that includes southeast coast
+    # of Singapore where there are hundreds of vessels.
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        response = client.post(
+            "/detections",
+            json={
+                "scene_id": "LC08_L1TP_125059_20240913_20240920_02_T1",
+                "scratch_path": tmp_dir,
+                "json_path": "",
+                "crop_path": "",
+            },
+        )
+        assert response.status_code == 200
+        predictions = response.json()["predictions"]
+        # There are many correct vessels in this scene.
+        assert len(predictions) >= 100