| |
| |
| """Check that python/eg2_image.py reproduces the upstream image processor bit for bit on the shipped test images: |
| the target size, the uint8 antialiased bicubic resize (vs torchvision) and the pixel values / patch layout (vs |
| transformers' Gemma4ImageProcessor, if transformers is installed). |
| |
| python3 check_preprocess.py [--images images] [--model-dir <google/embeddinggemma-2 snapshot>] |
| |
| Needs torch + torchvision (and transformers for the processor comparison). Exit code 0 = identical everywhere. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE.parent / "python")) |
| import eg2_image |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| ap.add_argument("--images", default=str(HERE / "images")) |
| ap.add_argument("--model-dir", default=None, help="checkpoint folder: also compare with Gemma4ImageProcessor") |
| args = ap.parse_args() |
| import torch |
| from torchvision.transforms.v2 import functional as tvF |
|
|
| proc = None |
| if args.model_dir: |
| from transformers import AutoProcessor |
|
|
| proc = AutoProcessor.from_pretrained(args.model_dir).image_processor |
| items = json.loads((Path(args.images) / "manifest.json").read_text()) |
| bad = 0 |
| for it in items: |
| rgb = eg2_image.load_rgb(Path(args.images) / it["file"]) |
| th, tw = eg2_image.target_size(rgb.shape[0], rgb.shape[1]) |
| mine = eg2_image.resize_bicubic_aa_u8(rgb, th, tw) |
| t = torch.from_numpy(rgb).permute(2, 0, 1).contiguous() |
| ref = tvF.resize(t, [th, tw], interpolation=tvF.InterpolationMode.BICUBIC, antialias=True).permute(1, 2, 0).numpy() if (th, tw) != rgb.shape[:2] else rgb |
| same = mine.shape == ref.shape and bool((mine == ref).all()) |
| msg = f"{it['file']:24s} {rgb.shape[1]}x{rgb.shape[0]} -> {tw}x{th}: resize {'identical' if same else 'DIFFERS'}" |
| if proc is not None: |
| feats = proc(images=[torch.from_numpy(rgb).permute(2, 0, 1)], return_tensors="pt") |
| pv, pos = feats["pixel_values"][0], feats["image_position_ids"][0] |
| vi = eg2_image.prepare_vision_inputs(Path(args.images) / it["file"], dtype=np.float32) |
| valid = pos[:, 0] >= 0 |
| key = {(int(x), int(y)): i for i, (x, y) in enumerate(pos[valid].tolist())} |
| perm = [key[(int(x), int(y))] for x, y in zip(vi.pos_x[0][: vi.n_soft * 9], vi.pos_y[0][: vi.n_soft * 9])] |
| same_pv = bool(np.array_equal(pv[perm].numpy(), vi.patches[0, : vi.n_soft * 9] * eg2_image.RESCALE)) and int(feats["num_soft_tokens_per_image"][0]) == vi.n_soft |
| same = same and same_pv |
| msg += f", processor pixel values {'identical' if same_pv else 'DIFFER'} (n_soft {vi.n_soft})" |
| bad += not same |
| print(msg) |
| print("CHECK_PREPROCESS", "PASS" if bad == 0 else f"FAIL ({bad} differ)") |
| return 0 if bad == 0 else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
