Skip to content

Commit

Permalink
feat(web_api): lift response size limit for GET cutout
Browse files Browse the repository at this point in the history
  • Loading branch information
nkemnitz committed Aug 30, 2024
1 parent 2e2db5a commit 0b316a9
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions web_api/app/painting.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
# pylint: disable=all # type: ignore
import gzip
from io import BytesIO
from typing import Annotated

import einops
import numpy as np
from fastapi import FastAPI, Query, Request
from fastapi.responses import Response
from fastapi.responses import StreamingResponse

from zetta_utils.geometry import Vec3D
from zetta_utils.layer.volumetric import VolumetricIndex
Expand Down Expand Up @@ -36,8 +37,16 @@ async def read_cutout(
if is_fortran:
data = einops.rearrange(data, "C X Y Z -> Z Y X C")
data_bytes = data.tobytes()
compressed_data = gzip.compress(data_bytes)
return Response(content=compressed_data)

def chunked_compress(data_bytes: bytes):
with BytesIO() as buffer:
with gzip.GzipFile(fileobj=buffer, mode='wb') as gzip_file:
gzip_file.write(data_bytes)
buffer.seek(0)
while chunk := buffer.read(64 * 1024):
yield chunk

return StreamingResponse(chunked_compress(data_bytes), media_type="application/gzip")


@api.post("/cutout")
Expand Down

0 comments on commit 0b316a9

Please sign in to comment.