guidellm.utils.vision

def encode_image(
    image: bytes | str | Path | np.ndarray | libs.Image,
    width: int | None = None,
    height: int | None = None,
    max_size: int | None = None,
    max_width: int | None = None,
    max_height: int | None = None,
    encode_type: Literal["base64", "url"] | None = "base64",
) -> dict[Literal["type", "image", "image_pixels", "image_bytes"], str | int | None]:
    """
    Input image types:
    - bytes: raw image bytes, decoded with Pillow
    - str: file path on disk, url, or already base64 encoded image string
    - pathlib.Path: file path on disk
    - np.ndarray: image array, decoded with Pillow
    - PIL.Image.Image: Pillow image
    - datasets.Image: HuggingFace datasets Image object

    max_size: maximum size of the longest edge of the image
    max_width: maximum width of the image
    max_height: maximum height of the image

    encode_type: None to return the supported format
        (url for url, base64 string for others)
        "base64" to return base64 encoded string (or download URL and encode)
        "url" to return url (only if input is url, otherwise fails)

    Returns a str of either:
    - image url
    - "data:image/{type};base64, {data}" string
    """
    if isinstance(image, str) and is_url(image):
        if encode_type == "base64":
            response = httpx.get(image)
            response.raise_for_status()
            return encode_image(
                image=response.content,
                max_size=max_size,
                max_width=max_width,
                max_height=max_height,
                encode_type="base64",
            )

        if any([width, height, max_size, max_width, max_height]):
            raise ValueError(f"Cannot resize image {image} when encode_type is 'url'")

        return {
            "type": "image_url",
            "image": image,
            "image_pixels": None,
            "image_bytes": None,
        }

    decoded_image: libs.Image

    if isinstance(image, bytes):
        decoded_image = libs.PILImage.open(io.BytesIO(image))
    elif isinstance(image, str) and image.startswith("data:image/"):
        _, encoded = image.split(",", 1)
        image_data = base64.b64decode(encoded)
        decoded_image = libs.PILImage.open(io.BytesIO(image_data))
    elif isinstance(image, str | Path):
        decoded_image = libs.PILImage.open(image)
    elif isinstance(image, np.ndarray):
        decoded_image = libs.PILImage.fromarray(image)
    elif isinstance(image, libs.Image):
        decoded_image = image
    else:
        raise ValueError(f"Unsupported image type: {type(image)} for {image}")

    output_image = resize_image(
        decoded_image,
        width=width,
        height=height,
        max_width=max_width,
        max_height=max_height,
        max_size=max_size,
    )
    if output_image.mode != "RGB":
        output_image = output_image.convert("RGB")

    buffer = io.BytesIO()
    output_image.save(buffer, format="JPEG")
    image_bytes = buffer.getvalue()
    image_base64 = base64.b64encode(image_bytes).decode("utf-8")

    return {
        "type": "image_base64",
        "image": f"data:image/jpeg;base64,{image_base64}",
        "image_pixels": output_image.width * output_image.height,
        "image_bytes": len(image_bytes),
    }

def encode_video(
    video: bytes | str | Path,
    encode_type: Literal["base64", "url"] | None = "base64",
) -> dict[
    Literal["type", "video", "video_frames", "video_seconds", "video_bytes"],
    str | int | float | None,
]:
    """
    Input video types:
    - bytes: raw video bytes
    - str: file path on disk, url, or already base64 encoded video string
    - pathlib.Path: file path on disk
    - datasets.Video: HuggingFace datasets Video object

    encode_type: None to return the supported format
        (url for url, base64 string for others)
        "base64" to return base64 encoded string (or download URL and encode)
        "url" to return url (only if input is url, otherwise fails)

    Returns a str of either:
    - video url
    - "data:video/{type};base64, {data}" string
    """
    if isinstance(video, str) and is_url(video):
        if encode_type == "base64":
            response = httpx.get(video)
            response.raise_for_status()
            return encode_video(video=response.content, encode_type="base64")

        return {
            "type": "video_url",
            "video": video,
            "video_frames": None,
            "video_seconds": None,
            "video_bytes": None,
        }

    if isinstance(video, str) and video.startswith("data:video/"):
        data_str = video.split(",", 1)[1]

        return {
            "type": "video_base64",
            "video": video,
            "video_frames": None,
            "video_seconds": None,
            "video_bytes": len(data_str) * 3 // 4,  # base64 to bytes
        }

    if isinstance(video, str | Path):
        path = Path(video)
        video_bytes = path.read_bytes()
        video_format = get_file_format(path)
    elif isinstance(video, bytes):
        video_bytes = video
        video_format = "unknown"
    else:
        raise ValueError(f"Unsupported video type: {type(video)} for {video}")

    video_base64 = base64.b64encode(video_bytes).decode("utf-8")

    return {
        "type": "video_base64",
        "video": f"data:video/{video_format};base64,{video_base64}",
        "video_frames": None,
        "video_seconds": None,
        "video_bytes": len(video_bytes),
    }

def get_file_format(path: Path | str) -> str:
    """Get file format from path extension."""
    suffix = Path(path).suffix.lower()
    return suffix[1:] if suffix.startswith(".") else "unknown"

def image_dict_to_pil(item: dict[str, Any]) -> libs.Image:
    """
    Decode an encoded image column item to a PIL Image for vLLM multi_modal_data.

    The item must have an "image" key with either a data URL (data:image/...;base64,...)
    or an HTTP(S) URL. For data URLs the image is base64-decoded; for URLs the
    image is fetched with httpx.

    :param item: Dict with "image" key (data URL or URL string)
    :return: PIL Image in RGB if needed
    :raises ValueError: If item has no "image" or unsupported format
    """
    image_spec = item.get("image")
    if not image_spec or not isinstance(image_spec, str):
        raise ValueError(
            "Encoded image item must have an 'image' key with a data URL or URL string."
        )
    if image_spec.startswith("data:image/"):
        _, encoded = image_spec.split(",", 1)
        data = base64.b64decode(encoded)
        decoded_image = libs.PILImage.open(io.BytesIO(data))
    elif image_spec.startswith(("http://", "https://")):
        response = httpx.get(image_spec)
        response.raise_for_status()
        decoded_image = libs.PILImage.open(io.BytesIO(response.content))
    else:
        raise ValueError(
            "Encoded image 'image' value must be a data:image/... URL or "
            f"http(s) URL, got: {image_spec[:80]!r}..."
        )
    if decoded_image.mode != "RGB":
        decoded_image = decoded_image.convert("RGB")  # type: ignore[assignment]
        # convert() returns Image; PILImage.open() may be ImageFile
    return decoded_image

def synthesize_image(
    width: int,
    height: int,
    *,
    content: str = "gradient",
    image_format: str = "jpeg",
    jpeg_quality: int = 85,
    seed: int = 0,
    row_index: int = 0,
) -> dict[str, Any]:
    """
    Synthesize a single image and return the canonical encoded dict.

    The output shape matches :func:`encode_image` so it flows through the rest
    of the pipeline (column mapper -> finalizer) unchanged.

    :param width: image width in pixels.
    :param height: image height in pixels.
    :param content: ``gradient`` (default, per-row randomized), ``noise``,
        ``solid``, or ``checkerboard``.
    :param image_format: ``jpeg`` (default) or ``png``.
    :param jpeg_quality: JPEG quality 1..100 (ignored for png).
    :param seed: base seed for reproducibility.
    :param row_index: row index used to vary the RNG stream per row so
        successive rows are byte-different even with the same seed.
    """
    if width <= 0 or height <= 0:
        raise ValueError(f"width and height must be positive, got {width}x{height}")
    fmt = image_format.lower()
    if fmt not in _SYNTHETIC_IMAGE_FORMATS:
        raise ValueError(
            f"Unsupported synthetic image format '{image_format}', "
            f"expected one of {_SYNTHETIC_IMAGE_FORMATS}"
        )
    if content not in _SYNTHETIC_IMAGE_CONTENT:
        raise ValueError(
            f"Unsupported synthetic image content '{content}', "
            f"expected one of {_SYNTHETIC_IMAGE_CONTENT}"
        )

    rng = _row_rng(seed, row_index)
    arr = _generate_image_array(height, width, content, rng)
    img = libs.PILImage.fromarray(arr, mode="RGB")

    buffer = io.BytesIO()
    if fmt == "jpeg":
        img.save(buffer, format="JPEG", quality=int(jpeg_quality), optimize=False)
        mime = "image/jpeg"
    else:
        img.save(buffer, format="PNG", optimize=False, compress_level=6)
        mime = "image/png"

    image_bytes = buffer.getvalue()
    image_b64 = base64.b64encode(image_bytes).decode("utf-8")

    return {
        "type": "image_base64",
        "image": f"data:{mime};base64,{image_b64}",
        "image_pixels": width * height,
        "image_bytes": len(image_bytes),
    }

def synthesize_video(
    width: int,
    height: int,
    frames: int,
    *,
    fps: float = 1.0,
    content: str = "gradient",
    video_format: str = "mp4",
    video_bitrate: str | None = None,
    seed: int = 0,
    row_index: int = 0,
) -> dict[str, Any]:
    """
    Synthesize a short video clip and return the canonical encoded dict.

    Matches the shape of :func:`encode_video`. Only ``mp4`` (h264, yuv420p) is
    supported in v1. Encoding uses ``-fflags +bitexact`` so two runs with the
    same seed produce byte-identical mp4 payloads.

    :param width: frame width in pixels (must be > 0).
    :param height: frame height in pixels (must be > 0).
    :param frames: number of frames in the clip (must be >= 1).
    :param fps: frames per second (encoded into the container, drives
        ``video_seconds = frames / fps``).
    :param content: ``gradient`` (default, per-frame randomized) or ``noise``.
    :param video_format: only ``mp4`` is supported in v1.
    :param video_bitrate: optional libx264 bitrate, e.g. ``"500k"``. ``None``
        leaves the codec at its default CRF-based rate control.
    :param seed: base seed for reproducibility.
    :param row_index: row index used to vary the RNG stream per row.
    """
    if width <= 0 or height <= 0:
        raise ValueError(f"width and height must be positive, got {width}x{height}")
    if frames <= 0:
        raise ValueError(f"frames must be positive, got {frames}")
    if fps <= 0:
        raise ValueError(f"fps must be positive, got {fps}")
    fmt = video_format.lower()
    if fmt not in _SYNTHETIC_VIDEO_FORMATS:
        raise ValueError(
            f"Unsupported synthetic video format '{video_format}', "
            f"expected one of {_SYNTHETIC_VIDEO_FORMATS}"
        )
    if content not in _SYNTHETIC_VIDEO_CONTENT:
        raise ValueError(
            f"Unsupported synthetic video content '{content}', "
            f"expected one of {_SYNTHETIC_VIDEO_CONTENT}"
        )

    rng = _row_rng(seed, row_index)
    clip = np.empty((frames, height, width, 3), dtype=np.uint8)
    for i in range(frames):
        frame_seed = int(rng.integers(0, 2**31 - 1))
        frame_rng = np.random.Generator(np.random.PCG64(frame_seed))
        if content == "gradient":
            clip[i] = _gradient_frame(height, width, frame_rng)
        else:
            clip[i] = _noise_frame(height, width, frame_rng)

    write_kwargs: dict[str, Any] = {
        "extension": ".mp4",
        "fps": float(fps),
        "codec": "libx264",
        "macro_block_size": 1,
        "ffmpeg_params": [
            "-fflags",
            "+bitexact",
            "-flags:v",
            "+bitexact",
        ],
    }
    if video_bitrate is not None:
        write_kwargs["bitrate"] = str(video_bitrate)

    video_bytes = libs.iio.imwrite("<bytes>", clip, **write_kwargs)  # type: ignore[attr-defined]
    video_b64 = base64.b64encode(video_bytes).decode("utf-8")

    return {
        "type": "video_base64",
        "video": f"data:video/mp4;base64,{video_b64}",
        "video_frames": int(frames),
        "video_seconds": float(frames) / float(fps),
        "video_bytes": len(video_bytes),
    }