Spaces:
Runtime error
Runtime error
File size: 6,045 Bytes
f4bd707 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
utils/predictor.py
Handles image preprocessing and running predictions through each model.
Returns: (label, confidence_score, input_tensor)
The input_tensor is passed back so Grad-CAM can use it.
"""
import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
# ──────────────────────────────────────────────────────────
# IMAGE PREPROCESSING
# ──────────────────────────────────────────────────────────
# Standard ImageNet normalization values — used because most
# deep learning models are pretrained on ImageNet.
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
def preprocess_image(image: Image.Image) -> torch.Tensor:
"""
Convert a PIL image into a normalized tensor ready for model input.
Steps:
1. Resize to 224x224 (standard input size for ResNet)
2. Convert to PyTorch tensor (values 0–1)
3. Normalize using ImageNet mean and std
Args:
image: A PIL Image (RGB).
Returns:
A tensor of shape [1, 3, 224, 224] — batch of 1 image.
"""
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
])
# Add a batch dimension: [3, 224, 224] → [1, 3, 224, 224]
tensor = transform(image).unsqueeze(0)
return tensor
# ──────────────────────────────────────────────────────────
# LABEL DEFINITIONS
# ──────────────────────────────────────────────────────────
TREE_LABELS = ["Non-Tree", "Tree"]
SPECIES_LABELS = ["Mango", "White Gum"]
STAGE_LABELS = ["Seedling", "Sapling", "Mature", "Overmature"]
# ──────────────────────────────────────────────────────────
# INFERENCE HELPER
# ──────────────────────────────────────────────────────────
def run_inference(model: torch.nn.Module, tensor: torch.Tensor) -> tuple:
"""
Run a forward pass through a model and return prediction results.
Args:
model: A loaded PyTorch model in eval mode.
tensor: Input tensor of shape [1, 3, 224, 224].
Returns:
(predicted_index, confidence_score)
"""
device = next(model.parameters()).device
tensor = tensor.to(device)
with torch.no_grad(): # No gradient tracking needed for inference
outputs = model(tensor) # Raw logits: [1, num_classes]
probs = torch.softmax(outputs, dim=1) # Convert to probabilities
confidence, predicted_idx = torch.max(probs, dim=1)
return predicted_idx.item(), confidence.item()
# ──────────────────────────────────────────────────────────
# PHASE 1: TREE vs NON-TREE
# ──────────────────────────────────────────────────────────
def predict_tree(image: Image.Image, model: torch.nn.Module) -> tuple:
"""
Predict whether the image contains a tree or not.
Args:
image: PIL Image.
model: Loaded tree vs non-tree model.
Returns:
(label: str, confidence: float, tensor: torch.Tensor)
e.g. ("Tree", 0.97, tensor)
"""
tensor = preprocess_image(image)
idx, confidence = run_inference(model, tensor)
label = TREE_LABELS[idx]
return label, confidence, tensor
# ──────────────────────────────────────────────────────────
# PHASE 2: SPECIES DETECTION
# ──────────────────────────────────────────────────────────
def predict_species(image: Image.Image, model: torch.nn.Module) -> tuple:
"""
Predict the tree species: Mango or White Gum.
Args:
image: PIL Image.
model: Loaded species detection model.
Returns:
(label: str, confidence: float, tensor: torch.Tensor)
e.g. ("Mango", 0.89, tensor)
"""
tensor = preprocess_image(image)
idx, confidence = run_inference(model, tensor)
label = SPECIES_LABELS[idx]
return label, confidence, tensor
# ──────────────────────────────────────────────────────────
# PHASE 3: STAGE CLASSIFICATION
# ──────────────────────────────────────────────────────────
def predict_stage(image: Image.Image, model: torch.nn.Module, species: str) -> tuple:
"""
Predict the growth stage of the tree.
Routes to the appropriate model based on detected species.
Args:
image: PIL Image.
model: The stage model for the detected species (mango or gum).
species: Species name (used for logging/display only here).
Returns:
(label: str, confidence: float, tensor: torch.Tensor)
e.g. ("Mature", 0.92, tensor)
"""
tensor = preprocess_image(image)
idx, confidence = run_inference(model, tensor)
label = STAGE_LABELS[idx]
return label, confidence, tensor
|