43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
def handle_reference_images(img_binary_api: Any = None, temp_file_ref: str = "", loaded_client_for_upload: Any = None, target_key: str = "first_image", **_: Any):
|
|
if temp_file_ref:
|
|
output = img_binary_api if isinstance(img_binary_api, list) else []
|
|
if hasattr(loaded_client_for_upload, "upload_file"):
|
|
output.append(loaded_client_for_upload.upload_file(temp_file_ref))
|
|
else:
|
|
output.append(temp_file_ref)
|
|
return output
|
|
|
|
if img_binary_api is not None and type(img_binary_api).__name__ == "Tensor":
|
|
img_array = (img_binary_api[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(img_array)
|
|
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
|
|
tmp_path = tmp.name
|
|
pil_image.save(tmp_path, format="PNG")
|
|
try:
|
|
uploaded = loaded_client_for_upload.upload_file(tmp_path) if hasattr(loaded_client_for_upload, "upload_file") else tmp_path
|
|
finally:
|
|
try:
|
|
os.unlink(tmp_path)
|
|
except OSError:
|
|
pass
|
|
return {target_key: uploaded}
|
|
|
|
if img_binary_api is not None and isinstance(img_binary_api, str) and img_binary_api:
|
|
if os.path.isfile(img_binary_api) and hasattr(loaded_client_for_upload, "upload_file"):
|
|
uploaded = loaded_client_for_upload.upload_file(img_binary_api)
|
|
else:
|
|
uploaded = img_binary_api
|
|
return {target_key: uploaded}
|
|
|
|
return img_binary_api if isinstance(img_binary_api, list) else []
|