50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
|
|
def tensor_to_pil_base64(image_tensor):
|
|
# Assuming the first dimension is the batch size, select the first image
|
|
if image_tensor.dim() > 3:
|
|
image_tensor = image_tensor[0] # Select the first image in the batch
|
|
|
|
# Ensure tensor is on the CPU
|
|
if image_tensor.is_cuda:
|
|
image_tensor = image_tensor.cpu()
|
|
|
|
# Check if it's a single color channel (grayscale) and needs color dimension expansion
|
|
if image_tensor.dim() == 2:
|
|
image_tensor = image_tensor.unsqueeze(0) # Add a channel dimension
|
|
|
|
# Convert from float tensors to uint8
|
|
if image_tensor.dtype == torch.float32:
|
|
image_tensor = (image_tensor * 255).byte()
|
|
|
|
# Permute the tensor dimensions if it's in C x H x W format to H x W x C for RGB
|
|
if image_tensor.dim() == 3 and image_tensor.size(0) == 3:
|
|
image_tensor = image_tensor.permute(1, 2, 0)
|
|
|
|
# Convert to numpy array
|
|
image_np = image_tensor.numpy()
|
|
|
|
# Convert numpy array to PIL Image
|
|
image_pil = Image.fromarray(image_np)
|
|
|
|
# Save the PIL image to a BytesIO buffer
|
|
buffer = BytesIO()
|
|
# You can change the format to PNG if you prefer
|
|
image_pil.save(buffer, format="PNG")
|
|
|
|
# Encode the buffer contents to base64
|
|
img_str = base64.b64encode(buffer.getvalue()).decode()
|
|
|
|
return img_str
|
|
|
|
|
|
# Example usage:
|
|
# Assuming `image_tensor` is your PyTorch tensor
|
|
# img_str = tensor_to_pil_base64(image_tensor)
|
|
# `img_str` is now a base64-encoded string of the image
|