| 12 | |
| 13 | |
| 14 | class VTP_Tokenizer: |
| 15 | def __init__( |
| 16 | self, |
| 17 | hf_model_path, |
| 18 | img_size=256, |
| 19 | horizon_flip=0.5, |
| 20 | fp16=True, |
| 21 | normalize_type="imagenet" |
| 22 | ): |
| 23 | """Initialize VTP Tokenizer. |
| 24 | |
| 25 | Args: |
| 26 | hf_model_path: Path to HuggingFace VTPModel directory |
| 27 | img_size: Input image size |
| 28 | horizon_flip: Horizontal flip probability for data augmentation |
| 29 | fp16: Whether to use FP16 precision |
| 30 | normalize_type: Normalization type, one of "half" (0.5 mean/std) or "imagenet" |
| 31 | """ |
| 32 | self.img_size = img_size |
| 33 | self.horizon_flip = horizon_flip |
| 34 | self.fp16 = fp16 |
| 35 | self.normalize_type = normalize_type |
| 36 | |
| 37 | # Setup normalization transforms |
| 38 | self._setup_normalization(normalize_type) |
| 39 | |
| 40 | # Load HuggingFace model |
| 41 | from vtp.models.vtp_hf import VTPModel |
| 42 | self.model = VTPModel.from_pretrained(hf_model_path) |
| 43 | self.model = self.model.cuda().eval() |
| 44 | |
| 45 | config = self.model.config |
| 46 | self.patch_size = config.vision_patch_size |
| 47 | self.embed_dim = config.vision_feature_bottleneck |
| 48 | |
| 49 | self.downsample_ratio = self.patch_size |
| 50 | self.latent_size = img_size // self.downsample_ratio |
| 51 | |
| 52 | print(f"VTP Tokenizer: patch_size={self.patch_size}, embed_dim={self.embed_dim}, " |
| 53 | f"downsample_ratio={self.downsample_ratio}, latent_size={self.latent_size}, " |
| 54 | f"normalize={self.normalize_type}") |
| 55 | |
| 56 | def _setup_normalization(self, normalize_type): |
| 57 | """Setup normalization and inverse normalization transforms.""" |
| 58 | if normalize_type == "half": |
| 59 | norm_cfg = NORMALIZE_HALF |
| 60 | elif normalize_type == "imagenet": |
| 61 | norm_cfg = NORMALIZE_IMAGENET |
| 62 | else: |
| 63 | raise ValueError(f"Unknown normalize_type: {normalize_type}. Use 'half' or 'imagenet'.") |
| 64 | |
| 65 | self.norm_mean = norm_cfg["mean"] |
| 66 | self.norm_std = norm_cfg["std"] |
| 67 | |
| 68 | # Inverse normalization: x_orig = x_norm * std + mean |
| 69 | # Which is: Normalize with mean=-mean/std, std=1/std |
| 70 | inv_mean = [-m / s for m, s in zip(self.norm_mean, self.norm_std)] |
| 71 | inv_std = [1.0 / s for s in self.norm_std] |