import torch import numpy as np from PIL import Image def convert(image): if isinstance(image, Image.Image): img_array = np.array(image) img_tensor = torch.from_numpy(img_array).float() / 255. if img_tensor.ndim == 3 and img_tensor.shape[-1] == 3: img_tensor = img_tensor.permute(2, 0, 1) image = img_tensor.unsqueeze(0).permute(0, 2, 3, 1) elif isinstance(image, torch.Tensor): img_array = image.squeeze(0).cpu().numpy() * 255.0 image = Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8)) return image