apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): correctly handle image stacks in persist stages

This commit is contained in:
Sean Sube 2023-11-27 21:45:32 -06:00
parent 828c9511a6
commit 042181b9c5
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/onnx_web/chain/persist_disk.py
View File

@ -31,7 +31,7 @@ class PersistDiskStage(BaseStage):
) -> StageResult: ) -> StageResult:
logger.info("persisting %s images to disk: %s", len(sources), output) logger.info("persisting %s images to disk: %s", len(sources), output)
for source, name in zip(sources, output): for source, name in zip(sources.as_image(), output):
dest = save_image(server, name, source, params=params, size=size) dest = save_image(server, name, source, params=params, size=size)
logger.info("saved image to %s", dest) logger.info("saved image to %s", dest)

10
api/onnx_web/chain/persist_s3.py
View File

@ -1,6 +1,6 @@
from io import BytesIO from io import BytesIO
from logging import getLogger from logging import getLogger
from typing import Optional from typing import List, Optional
from boto3 import Session from boto3 import Session
from PIL import Image from PIL import Image
@ -23,7 +23,7 @@ class PersistS3Stage(BaseStage):
_params: ImageParams, _params: ImageParams,
sources: StageResult, sources: StageResult,
*, *,
output: str, output: List[str],
bucket: str, bucket: str,
endpoint_url: Optional[str] = None, endpoint_url: Optional[str] = None,
profile_name: Optional[str] = None, profile_name: Optional[str] = None,
@ -33,14 +33,14 @@ class PersistS3Stage(BaseStage):
session = Session(profile_name=profile_name) session = Session(profile_name=profile_name)
s3 = session.client("s3", endpoint_url=endpoint_url) s3 = session.client("s3", endpoint_url=endpoint_url)
for source in sources.as_image(): for source, name in zip(sources.as_image(), output):
data = BytesIO() data = BytesIO()
source.save(data, format=server.image_format) source.save(data, format=server.image_format)
data.seek(0) data.seek(0)
try: try:
s3.upload_fileobj(data, bucket, output) s3.upload_fileobj(data, bucket, name)
logger.info("saved image to s3://%s/%s", bucket, output) logger.info("saved image to s3://%s/%s", bucket, name)
except Exception: except Exception:
logger.exception("error saving image to S3") logger.exception("error saving image to S3")