File size: 1,073 Bytes
28dec30 e81fa57 28dec30 e81fa57 28dec30 e81fa57 28dec30 e81fa57 28dec30 e81fa57 28dec30 e81fa57 28dec30 e81fa57 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import io
from PIL import Image
import numpy as np
import torch
class ImageEncoder:
@torch.inference_mode()
def encode_torch(self, img: torch.Tensor, quality=95):
if img.ndim == 2:
img = (
img[None]
.repeat_interleave(3, dim=0)
.permute(1, 2, 0)
.contiguous()
.clamp(0, 255)
.type(torch.uint8)
)
elif img.ndim == 3:
if img.shape[0] == 3:
img = img.permute(1, 2, 0).contiguous().clamp(0, 255).type(torch.uint8)
elif img.shape[2] == 3:
img = img.contiguous().clamp(0, 255).type(torch.uint8)
else:
raise ValueError(f"Unsupported image shape: {img.shape}")
else:
raise ValueError(f"Unsupported image num dims: {img.ndim}")
img = img.cpu().numpy().astype(np.uint8)
im = Image.fromarray(img)
iob = io.BytesIO()
im.save(iob, format="JPEG", quality=quality)
iob.seek(0)
return iob
|