| """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() |
|
|
