"""Clean crop-level ripeness inference for the plantjepa-preview release. Self-contained on purpose: the exact same file runs locally and ships to HuggingFace. Only dependencies are torch, Pillow, and numpy. The embedded encoder and crop are faithful copies of the training-time code (models/vit.py SmallViT, data/ripeness.py crop_instance); see tests/test_inference_parity.py. The checkpoint bundle must contain: encoder_config.backbone, encoder_state_dict, scaler_mean, scaler_scale, classifier_coef, classifier_intercept, resolution, classes, ordered_stages. Usage: from inference import RipenessModel model = RipenessModel.load("jepa_ripeness_model.pt", "calibration.json") print(model.predict("image.jpg", bbox=(x1, y1, x2, y2))) """ import argparse import json import math import numpy as np import torch from PIL import Image, ImageOps class _SmallViT(torch.nn.Module): def __init__(self, resolution=224, patch_size=16, dim=192, depth=4, heads=3): super().__init__() if resolution % patch_size: raise ValueError('resolution must be divisible by patch_size') self.patch = torch.nn.Conv2d(3, dim, patch_size, patch_size) self.position = torch.nn.Parameter(torch.randn(1, (resolution // patch_size) ** 2, dim) * .02) layer = torch.nn.TransformerEncoderLayer(dim, heads, dim * 4, dropout=0., batch_first=True, activation='gelu', norm_first=True) self.blocks = torch.nn.TransformerEncoder(layer, depth, enable_nested_tensor=False) self.norm = torch.nn.LayerNorm(dim) self.dim = dim def encode(self, frames): shape = frames.shape x = frames.reshape(-1, *shape[-3:]) x = self.patch(x).flatten(2).transpose(1, 2) x = self.norm(self.blocks(x + self.position)).mean(1) return x.reshape(*shape[:-3], self.dim) forward = encode def _crop_fruit(image, bbox, resolution, padding=0.15): """Crop a fruit box with proportional context and square-pad it.""" image = ImageOps.exif_transpose(image).convert("RGB") x1, y1, x2, y2 = [float(v) for v in bbox] w, h = image.size if x2 > w or y2 > h or x2 <= x1 or y2 <= y1: raise ValueError(f"Bounding box out of range for {w}x{h} image: {bbox}") dx, dy = (x2 - x1) * padding, (y2 - y1) * padding box = (max(0, int(x1 - dx)), max(0, int(y1 - dy)), min(w, math.ceil(x2 + dx)), min(h, math.ceil(y2 + dy))) return ImageOps.pad(image.crop(box), (resolution, resolution), method=Image.Resampling.BICUBIC) def _softmax(logits): shifted = logits - logits.max(axis=-1, keepdims=True) exp = np.exp(shifted) return exp / exp.sum(axis=-1, keepdims=True) class RipenessModel: """JEPA encoder + linear ripeness head with optional temperature scaling.""" def __init__(self, bundle, temperature=1.0, device="cpu"): temperature = float(temperature) if temperature <= 0: raise ValueError("Temperature must be positive") missing = {"encoder_config", "encoder_state_dict", "scaler_mean", "scaler_scale", "classifier_coef", "classifier_intercept", "resolution", "classes", "ordered_stages"} if missing - set(bundle): raise ValueError(f"Checkpoint bundle missing keys: {sorted(missing - set(bundle))}") self.device = torch.device(device) self.encoder = _SmallViT(**bundle["encoder_config"]["backbone"]) self.encoder.load_state_dict(bundle["encoder_state_dict"], strict=True) self.encoder.eval().to(self.device) self.mean = bundle["scaler_mean"].to(self.device) self.scale = bundle["scaler_scale"].to(self.device) self.coef = bundle["classifier_coef"].to(self.device) self.intercept = bundle["classifier_intercept"].to(self.device) self.resolution = int(bundle["resolution"]) self.classes = list(bundle["classes"]) self.stages = list(bundle["ordered_stages"]) self.temperature = temperature @classmethod def load(cls, checkpoint_path, calibration_path=None, temperature=1.0, device="cpu"): if calibration_path is not None: with open(calibration_path) as handle: temperature = json.load(handle)["temperature"] bundle = torch.load(checkpoint_path, map_location="cpu", weights_only=True) return cls(bundle, temperature=temperature, device=device) def predict(self, image_path, bbox): with Image.open(image_path) as image: return self.predict_image(image, bbox) def predict_image(self, image, bbox): crop = _crop_fruit(image, bbox, self.resolution) frame = torch.from_numpy(np.asarray(crop).copy()).permute(2, 0, 1).float().div_(255).unsqueeze(0) with torch.inference_mode(): feature = self.encoder.encode(frame.to(self.device)).float() standardized = (feature - self.mean) / self.scale logits = torch.nn.functional.linear(standardized, self.coef, self.intercept) probs = _softmax(logits.cpu().numpy() / self.temperature)[0] stage_probs = {stage: float(probs[self.classes.index(stage)]) if stage in self.classes else 0.0 for stage in self.stages} predicted = max(stage_probs, key=stage_probs.get) maturity = 100.0 * sum(rank * value for rank, value in enumerate(stage_probs.values())) / (len(self.stages) - 1) return {"predicted_stage": predicted, "confidence": float(stage_probs[predicted]), "stage_probabilities": stage_probs, "ordinal_maturity_index_0_100": float(maturity), "temperature": self.temperature, "note": "Ordinal maturity index is an expected rank, not a calibrated physical measurement."} def main(): parser = argparse.ArgumentParser(description="Ripeness inference with a packaged JEPA checkpoint") parser.add_argument("--checkpoint", required=True) parser.add_argument("--image", required=True) parser.add_argument("--bbox", type=float, nargs=4, required=True, metavar=("X1", "Y1", "X2", "Y2")) parser.add_argument("--calibration", default=None) parser.add_argument("--temperature", type=float, default=1.0) parser.add_argument("--device", default="cpu") args = parser.parse_args() model = RipenessModel.load(args.checkpoint, args.calibration, args.temperature, args.device) print(json.dumps(model.predict(args.image, tuple(args.bbox)), indent=2)) if __name__ == "__main__": main()