Spaces:
Configuration error
Configuration error
File size: 1,850 Bytes
3070946 254b8d8 3070946 254b8d8 3070946 254b8d8 3070946 254b8d8 3070946 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Image preprocessing pipeline for disease and soil CNN models."""
import io
import hashlib
from PIL import Image
import numpy as np
try:
import torch
from torchvision import transforms
inference_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
except ImportError:
torch = None
inference_transform = None
def preprocess_image(image_bytes: bytes) -> tuple:
"""
Validate, preprocess and return a (1, 3, 224, 224) float32 tensor/ndarray
plus a deterministic cache key (SHA-256 of raw bytes).
Raises ValueError for invalid / non-image data.
"""
# --- validate ---
try:
img = Image.open(io.BytesIO(image_bytes))
img.verify() # catches truncated files
img = Image.open(io.BytesIO(image_bytes)) # reopen after verify
except Exception as exc:
raise ValueError(f"Invalid image file: {exc}") from exc
# --- convert to RGB (handles RGBA, palette, grayscale) ---
img = img.convert("RGB")
# --- reject suspiciously small images ---
w, h = img.size
if w < 64 or h < 64:
raise ValueError(f"Image too small ({w}×{h}). Minimum 64×64 px required.")
# --- transform ---
if torch is not None and inference_transform is not None:
tensor = inference_transform(img).unsqueeze(0) # (1, 3, 224, 224)
else:
resized = img.resize((224, 224))
arr = np.array(resized, dtype=np.float32) / 255.0
tensor = np.transpose(arr, (2, 0, 1))[np.newaxis, ...]
# --- cache key ---
cache_key = hashlib.sha256(image_bytes).hexdigest()
return tensor, cache_key
def bytes_to_pil(image_bytes: bytes) -> Image.Image:
"""Return a PIL Image from raw bytes, converted to RGB."""
return Image.open(io.BytesIO(image_bytes)).convert("RGB")
|