File size: 3,882 Bytes
0efe9d5 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """Quicksviewer image ingress: PIL bicubic RGB, CPU uint8, no normalization."""
import base64
import io
import os
from urllib.request import urlopen
import numpy as np
import torch
from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor
from transformers.feature_extraction_utils import BatchFeature
SIGLIP_SIZE = 980
def read_source(source):
if isinstance(source, os.PathLike):
source = os.fspath(source)
if isinstance(source, str) and source.startswith("data:"):
header, payload = source.split(",", 1)
if ";base64" not in header:
raise ValueError("Media data URLs must use base64 encoding")
return io.BytesIO(base64.b64decode(payload))
if isinstance(source, str) and source.startswith(("https://", "http://")):
with urlopen(source, timeout=120) as response:
return io.BytesIO(response.read())
if isinstance(source, str) and source.startswith("file://"):
return source[7:]
if isinstance(source, bytes):
return io.BytesIO(source)
return source
def to_pil(image):
if isinstance(image, Image.Image):
return image
if isinstance(image, (str, bytes, os.PathLike)) or hasattr(image, "read"):
with Image.open(read_source(image)) as opened:
return opened.convert("RGB")
if isinstance(image, torch.Tensor):
image = image.detach().cpu().numpy()
if not isinstance(image, np.ndarray):
raise TypeError(f"Unsupported image type: {type(image).__name__}")
if image.ndim == 3 and image.shape[0] in (1, 3, 4) and image.shape[-1] not in (1, 3, 4):
image = np.moveaxis(image, 0, -1)
if image.ndim not in (2, 3):
raise ValueError("An image must have shape HWC or CHW")
if image.dtype != np.uint8:
if not np.isfinite(image).all() or image.min() < 0 or image.max() > 255:
raise ValueError("Image values must be finite and in [0,255] (or [0,1])")
if np.issubdtype(image.dtype, np.floating) and image.max() <= 1:
image = image * 255
image = np.rint(image).astype(np.uint8)
if image.ndim == 3 and image.shape[-1] == 1:
image = image[..., 0]
return Image.fromarray(image)
class QuicksviewerImageProcessor(BaseImageProcessor):
model_input_names = ["pixel_values"]
def __init__(self, size=None, **kwargs):
super().__init__(**kwargs)
self.size = size or {"height": SIGLIP_SIZE, "width": SIGLIP_SIZE}
if self.size != {"height": SIGLIP_SIZE, "width": SIGLIP_SIZE}:
raise ValueError("Quicksviewer requires 980x980 images")
self.do_normalize = False
self.do_rescale = False
def preprocess_one(self, image):
image = to_pil(image).convert("RGB").resize(
(SIGLIP_SIZE, SIGLIP_SIZE), Image.Resampling.BICUBIC
)
array = np.array(image, dtype=np.uint8, copy=True)
return torch.from_numpy(array).permute(2, 0, 1).contiguous()
def preprocess(self, images, return_tensors="pt", **kwargs):
if kwargs:
raise TypeError(f"Unsupported image options: {sorted(kwargs)}")
if isinstance(images, (np.ndarray, torch.Tensor)) and images.ndim == 4:
images = list(images)
elif not isinstance(images, (list, tuple)):
images = [images]
frames = [self.preprocess_one(image) for image in images]
pixels = torch.stack(frames) if frames else torch.empty(
(0, 3, SIGLIP_SIZE, SIGLIP_SIZE), dtype=torch.uint8
)
if return_tensors not in (None, "pt", "np"):
raise ValueError("return_tensors must be 'pt', 'np', or None")
if return_tensors == "np":
pixels = pixels.numpy()
return BatchFeature({"pixel_values": pixels})
QuicksviewerImageProcessor.register_for_auto_class("AutoImageProcessor")
|