nnU-Net API¶

@dataclass
class nnUNetAPI:
    """
    General nnU-Net API.

    The ``writing_process_pool`` is only implemented for the
    command line entrypoints and, as such, is only used when ``app`` is None.

    Args:
        app: FastAPI application
        writing_process_pool: ProcessPool for writing files.
    """

    app: fastapi.FastAPI | None = None
    writing_process_pool: ProcessPool | None = None

    def __post_init__(self):
        """
        Initializes the nnUNetAPI instance, checking for GPU and setting up the DB.
        """
        if torch.cuda.is_available() is False:
            raise ValueError("No GPU available")
        self.model_dictionary, self.alias_dict = get_model_dictionary()
        # Initialise SQLite DB for zip storage
        self._db_path = Path(f"{NNUNET_OUTPUT_DIR}/zip_store.db")
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._db_conn = sqlite3.connect(self._db_path)
        self._init_db()

    def _init_db(self) -> None:
        """
        Creates the zip_store table if it does not exist.

        Schema:
            job_id (TEXT): Primary key.
            created_at (DATE): ISO date of creation.
            zip_path (TEXT): Path to the generated zip file.
        """
        cur = self._db_conn.cursor()
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS zip_store (
                job_id TEXT PRIMARY KEY,
                created_at DATE,
                zip_path TEXT
            )
            """
        )
        self._db_conn.commit()

    def _store_zip(self, job_id: str, zip_path: Path) -> None:
        """Inserts or replaces a record for a generated zip file.

        Args:
            job_id (str): Unique job identifier.
            zip_path (Path): Path to the generated zip file.
        """
        cur = self._db_conn.cursor()
        cur.execute(
            "INSERT OR REPLACE INTO zip_store (job_id, created_at, zip_path) VALUES (?, ?, ?)",
            (job_id, datetime.date.today().isoformat(), str(zip_path)),
        )
        self._db_conn.commit()

    def _get_zip_path(self, job_id: str) -> Path | None:
        """
        Retrieves the zip path for a given job_id.

        Args:
            job_id (str): Unique job identifier.

        Returns:
            Path | None: The path to the zip file if found, else None.
        """
        cur = self._db_conn.cursor()
        cur.execute(
            "SELECT zip_path FROM zip_store WHERE job_id = ?", (job_id,)
        )
        row = cur.fetchone()
        return Path(row[0]) if row else None

    def cleanup_old_records(self, days: int = 7) -> int:
        """
        Delete records older than *days*.
        Returns the number of rows removed.

        Args:
            days (int, optional): Number of days to keep records. Defaults to 7.
        """
        cutoff = datetime.date.today() - datetime.timedelta(days=days)
        cur = self._db_conn.cursor()
        cur.execute(
            "DELETE FROM zip_store WHERE created_at < ?", (cutoff.isoformat(),)
        )
        removed = cur.rowcount
        self._db_conn.commit()
        return removed

    def init_api(self):
        """
        Initializes the API.
        """
        if self.app is None:
            raise ValueError("app must be defined before init_api is called")
        self.app.add_api_route(
            "/infer",
            self.infer,
            methods=["POST"],
            response_model=InferenceResponse,
        )
        self.app.add_api_route(
            "/infer_file",
            self.infer_file,
            methods=["POST"],
            response_model=InferenceFileResponse,
        )
        self.app.add_api_route(
            "/infer_orthanc",
            self.infer_orthanc,
            methods=["POST"],
            response_model=InferenceResponse,
        )
        self.app.add_api_route(
            "/download/{job_id}",
            self.download_file,
            methods=["GET"],
            response_class=FileResponse,
        )
        self.app.add_api_route(
            "/model_info",
            self.model_info,
            methods=["GET"],
            response_model=ModelInfoResponse,
        )
        self.app.add_api_route(
            "/model_info_clean",
            self.model_info_clean,
            methods=["GET"],
            response_model=ModelInfoResponse,
        )
        self.app.add_api_route(
            "/request-params",
            self.request_params,
            methods=["GET"],
            response_model=JSONSchema,
        )
        self.app.add_api_route(
            "/healthz",
            self.healthz,
            methods=["GET"],
            response_model=HealthzResponse,
        )
        self.app.add_api_route(
            "/readyz",
            self.readyz,
            methods=["GET"],
            response_model=ReadyzResponse,
        )
        self.app.add_api_route(
            "/expire",
            self.expire,
            methods=["GET"],
            response_model=ExpireResponse,
        )

    def expire(self):
        """
        Calls the TTL cache expire method.
        """
        n_items = 0
        try:
            n_items = len(CACHE.expire())
        except Exception as e:
            logger.error("Failed to expire cache: %s", e)
            return JSONResponse(
                status_code=500,
                content={"status": "error", "message": str(e)},
            )
        return JSONResponse(
            status_code=200,
            content={"status": "ok", "message": f"Expired {n_items} items"},
        )

    def model_info(self):
        """
        Returns the model information.

        Returns:
            dict: Model information.
        """
        return self.model_info_clean()

    def model_info_clean(self):
        """
        Returns the model information with cleaned metadata.

        Returns:
            dict: Model information.
        """
        model_dict = deepcopy(self.model_dictionary)
        for model in model_dict.values():
            sd = SegWriter.init_from_metadata_dict(
                model["metadata"]
            ).segment_descriptions
            model_labels = {
                v: k for k, v in model["model_information"]["labels"].items()
            }
            for i in range(len(sd)):
                try:
                    label = sd[i][0x0062, 0x0005].value
                except KeyError:
                    label = None
                try:
                    meaning = sd[i][0x0062, 0x000F][0][0x0008, 0x0104].value
                except KeyError:
                    meaning = None
                try:
                    laterality = sd[i][0x0062, 0x0011][0][0x0008, 0x0104].value
                except KeyError:
                    laterality = None
                sd[i] = {
                    "Label ID": model_labels[i + 1],
                    "Name": label,
                    "Index": i + 1,
                }

            model["metadata"] = sd
            model["description"] = "\n".join(
                [
                    "Segments the following regions:",
                    ", ".join([model_labels[i + 1] for i in range(len(sd))]),
                    "Number of input channels:",
                    str(len(model["model_information"]["channel_names"])),
                ]
            )
            model["description_long"] = "\n".join(
                [
                    "Segments the following regions:",
                    *["\t- " + dict_to_str(sd[i]) for i in range(len(sd))],
                    "Uses the following channels:",
                    "\t- "
                    + dict_to_str(model["model_information"]["channel_names"]),
                ]
            )
        return model_dict

    def request_params(self):
        """
        Returns the request parameters.

        Returns:
            dict: Request parameters.
        """
        return InferenceRequest.model_json_schema()

    def _failure_payload(
        self,
        error: str,
        request_payload: dict | None,
        **extra_fields,
    ) -> dict[str, Any]:
        """
        Build a standardized failure payload for API responses.

        Args:
            error (str): Error message.
            request_payload (dict | None): Request payload.
            **extra_fields: Extra fields to add to the payload.

        Returns:
            dict[str, Any]: Failure payload.
        """
        payload = {
            "time_elapsed": None,
            "gpu": None,
            "nnunet_path": None,
            "metadata": None,
            "status": FAILURE_STATUS,
            "request": request_payload if request_payload is not None else {},
            "error": error,
        }
        payload.update(extra_fields)
        return payload

    def _raise_or_error_response(
        self,
        error: str,
        status_code: int,
        request_payload: dict | None,
        exception_type: type[Exception] = ValueError,
        **extra_fields,
    ):
        """
        Raise in CLI mode or return standardized JSON error in API mode.

        Args:
            error (str): Error message.
            status_code (int): status code to be used for error response.
            request_payload (dict | None): payload for the request.
            exception_type (type[Exception] = ValueError): type of Exception.
            **extra_fields: Extra fields to add to the payload.

        Returns:
            JSONResponse if self.app is not None.

        Raises:
            ``exception_type`` error.
        """
        if self.app is None:
            raise exception_type(error)
        return JSONResponse(
            content=self._failure_payload(
                error=error,
                request_payload=request_payload,
                **extra_fields,
            ),
            status_code=status_code,
        )

    async def healthz(self):
        """
        Returns a simple health check.
        """
        return {"status": "ok"}

    async def readyz(self):
        """
        Returns a readiness check.
        """
        models_loaded = len(self.model_dictionary) > 0
        gpu_available = False
        max_free_mem = None
        try:
            if torch.cuda.is_available():
                mem = get_gpu_memory()
                gpu_available = len(mem) > 0
                max_free_mem = max(mem) if mem else None
        except Exception:
            gpu_available = False
        status = (
            "ok"
            if models_loaded
            and (gpu_available or not torch.cuda.is_available())
            else "starting"
        )
        return {
            "status": status,
            "models_loaded": models_loaded,
            "gpu_available": gpu_available,
            "max_free_mem": max_free_mem,
        }

    async def infer(
        self, inference_request: Annotated[InferenceRequest, Query()]
    ):
        """
        Performs inference.

        Args:
            inference_request (InferenceRequest): Inference request.

        Returns:
            JSONResponse: Inference response.
        """
        if self.app is not None and self.writing_process_pool is not None:
            raise ValueError("Cannot use both app and writing_process_pool")
        params = normalize_inference_params(inference_request)
        add_file_handler_to_manager(
            log_path=os.path.join(params["output_dir"], "nnunet_serve.log"),
            exclude=[
                "nnunet_serve.entrypoints.entrypoint_batch",
                "nnunet_serve.entrypoints.entrypoint",
                "nnunet_serve.process_pool",
                "nnunet_serve.seg_writers",
            ],
        )
        nnunet_id = params["nnunet_id"]
        if isinstance(nnunet_id, str):
            nnunet_id = [nnunet_id]

        initial_default_args = resolve_models(
            nnunet_id=nnunet_id,
            model_dictionary=self.model_dictionary,
            alias_dict=self.alias_dict,
        )[2]
        may_inject = get_may_inject(initial_default_args)

        nnunet_id, insert_at = expand_cascade_inputs(
            params=params,
            nnunet_id=nnunet_id,
            model_dictionary=self.model_dictionary,
            alias_dict=self.alias_dict,
            may_inject_series=may_inject,
        )

        (
            nnunet_path,
            metadata,
            default_args,
            is_totalseg,
            min_mem,
            model_resolution_error,
        ) = resolve_models(
            nnunet_id=nnunet_id,
            model_dictionary=self.model_dictionary,
            alias_dict=self.alias_dict,
        )
        if model_resolution_error is not None:
            return self._raise_or_error_response(
                error=model_resolution_error,
                status_code=404,
                request_payload=params,
            )

        apply_request_defaults(
            params=params,
            default_args=default_args,
            inference_request=inference_request,
            insert_at=insert_at,
        )
        params["min_mem"] = min_mem

        if params.get("save_proba_map", False) and all(
            [x is None for x in params.get("proba_threshold", [])]
        ):
            error_str = (
                "proba_threshold must be not-None if save_proba_map is True"
            )
            return self._raise_or_error_response(
                error=error_str,
                status_code=400,
                request_payload=params,
            )

        series_paths, code, error_msg = get_series_paths(
            params["study_path"],
            series_folders=params["series_folders"],
            n=len(nnunet_id) if isinstance(nnunet_id, list) else None,
        )

        if code == FAILURE_STATUS:
            error_str = error_msg
            return self._raise_or_error_response(
                error=error_str,
                status_code=400,
                request_payload=params,
            )

        try:
            wait_for_gpu(min_mem)
        except (RuntimeError, TimeoutError) as e:
            error_str = str(e)
            return self._raise_or_error_response(
                error=error_str,
                status_code=503,
                request_payload=params,
                exception_type=RuntimeError,
            )

        if "tta" in params:
            mirroring = params["tta"]
        else:
            mirroring = True

        a = time.time()
        (
            output_paths,
            identifiers,
            is_empty,
            status,
            error,
        ) = run_predict_inference(
            series_paths=series_paths,
            metadata=metadata,
            mirroring=mirroring,
            params=params,
            nnunet_path=nnunet_path,
            is_totalseg=is_totalseg,
            writing_process_pool=self.writing_process_pool,
        )
        b = time.time()
        if status == FAILURE_STATUS:
            error_str = error
            return self._raise_or_error_response(
                error=error_str,
                status_code=500,
                request_payload=params,
                identifiers=identifiers,
                is_empty=is_empty,
                **output_paths,
            )
        payload = build_infer_success_payload(
            time_elapsed=b - a,
            nnunet_path=nnunet_path,
            metadata=metadata,
            request_params=params,
            identifiers=identifiers,
            is_empty=is_empty,
            output_paths=output_paths,
        )
        return JSONResponse(content=payload, status_code=200)

    async def infer_file(
        self,
        inference_request: Request,
        file: UploadFile = File(...),
    ):
        """
        Accept a file (or archive) upload, stores it, builds an InferenceRequest,
        and delegates to the existing ``infer`` method.
        """

        job_id = uuid.uuid4().hex
        form = await inference_request.form()
        json_str = form.get("request")
        if json_str is not None:
            payload = json.loads(json_str)
        else:
            payload = await inference_request.json()

        study_path = get_study_path(job_id)
        payload["study_path"] = str(study_path / "inputs")
        payload["output_dir"] = str(study_path / "output")

        try:
            inference_req = InferenceRequest(**payload)
        except Exception as exc:
            return fastapi.responses.JSONResponse(
                content=self._failure_payload(
                    error=f"Invalid request payload: {exc}",
                    request_payload=payload,
                ),
                status_code=422,
            )

        try:
            store_uploaded_file(file, job_id=job_id)
        except Exception as exc:
            return fastapi.responses.JSONResponse(
                content=self._failure_payload(
                    error=f"Failed to store uploaded file: {exc}",
                    request_payload={},
                ),
                status_code=400,
            )

        response = await self.infer(inference_req)

        if response.status_code == 200:
            zip_path = zip_directory(Path(inference_req.output_dir))
            self._store_zip(job_id, zip_path)
            shutil.rmtree(inference_req.study_path)
            shutil.rmtree(inference_req.output_dir)
            original = json.loads(response.body)
            original.update({"job_id": job_id})
            return JSONResponse(content=original, status_code=200)
        else:
            error_payload = json.loads(response.body)
            error_payload.update({"job_id": job_id})
            return JSONResponse(
                content=error_payload, status_code=response.status_code
            )

    async def infer_orthanc(
        self, inference_request: Annotated[InferenceRequestOrthanc, Query()]
    ):
        """Run inference for Orthanc-backed inputs and push SEG back to Orthanc.

        This adapter:
        1. Downloads Orthanc series referenced in ``series_ids``.
        2. Rewrites ``series_folders`` to local downloaded paths.
        3. Reuses ``infer``.
        4. Uploads produced DICOM SEG files back to Orthanc.

        Entries using ``from:`` are preserved as-is to keep cascade behavior.
        """
        job_id = uuid.uuid4().hex
        study_path = get_study_path(job_id)
        inputs_path = study_path / "inputs"
        output_path = study_path / "output"
        inputs_path.mkdir(parents=True, exist_ok=True)

        payload = inference_request.model_dump()
        payload["study_path"] = str(inputs_path)
        payload["output_dir"] = str(output_path)

        raw_series_ids = payload.get("series_ids", None)
        if raw_series_ids is None:
            return self._raise_or_error_response(
                error="series_ids must be defined",
                status_code=400,
                request_payload=payload,
            )

        if isinstance(raw_series_ids, list) and (
            len(raw_series_ids) == 0 or isinstance(raw_series_ids[0], str)
        ):
            series_ids = [raw_series_ids]
        else:
            series_ids = raw_series_ids

        orthanc_series_ids = []
        for stage_series in series_ids:
            for sid in stage_series:
                if isinstance(sid, str) and sid.startswith("from:"):
                    continue
                orthanc_series_ids.append(sid)

        local_series_map = {}
        if len(orthanc_series_ids) > 0:
            unique_series_ids = sorted(set(orthanc_series_ids))
            downloaded_paths = download_series(
                unique_series_ids, output_dir=str(inputs_path)
            )
            for sid, folder_path in downloaded_paths.items():
                local_series_map[sid] = os.path.relpath(
                    folder_path, inputs_path
                )

        adapted_series_folders = []
        for stage_series in series_ids:
            adapted_stage = []
            for sid in stage_series:
                if isinstance(sid, str) and sid.startswith("from:"):
                    adapted_stage.append(sid)
                else:
                    adapted_stage.append(local_series_map[sid])
            adapted_series_folders.append(adapted_stage)

        payload.pop("series_ids", None)
        payload["series_folders"] = adapted_series_folders

        try:
            infer_request = InferenceRequest(**payload)
        except Exception as exc:
            shutil.rmtree(study_path, ignore_errors=True)
            return self._raise_or_error_response(
                error=f"Invalid adapted request payload: {exc}",
                status_code=422,
                request_payload=payload,
            )

        response = await self.infer(infer_request)
        if response.status_code != 200:
            shutil.rmtree(study_path, ignore_errors=True)
            return response

        response_payload = json.loads(response.body)
        dicom_seg_paths = response_payload.get("dicom_segmentation", [])
        dicom_seg_paths = [p for p in dicom_seg_paths if p is not None]

        uploaded_instances = []
        if len(dicom_seg_paths) > 0:
            uploaded_instances = upload_series(dicom_seg_paths)
        response_payload["orthanc_upload"] = {
            "uploaded_instance_count": len(uploaded_instances),
            "responses": uploaded_instances,
        }

        shutil.rmtree(study_path, ignore_errors=True)
        return JSONResponse(content=response_payload, status_code=200)

    async def download_file(self, job_id: str):
        """
        Serve the zip file created for ``job_id`` (``job_id`` is the value returned by
        the ``infer_file`` endpoint).
        Returns 404 if the ``job_id`` is unknown or the file has been cleaned up.
        """
        zip_path = self._get_zip_path(job_id)
        if not zip_path or not zip_path.exists():
            raise fastapi.HTTPException(
                status_code=404, detail="Zip not found for job_id"
            )
        return FileResponse(
            path=zip_path, media_type="application/zip", filename=zip_path.name
        )

    def __del__(self):
        if hasattr(self, "_db_conn"):
            self._db_conn.close()