infer.py

#!/usr/bin/env python3
"""infer.py - simple CLI for running inference with birds_mobilenetv2.keras

Usage:
  python infer.py --image PATH [--model MODEL] [--labels LABELS] [--topk N]
"""
import argparse
import os
import sys
from PIL import Image
import numpy as np

try:
    from tensorflow.keras.models import load_model
except Exception as e:
    print('Failed to import TensorFlow Keras:', e, file=sys.stderr)
    raise


def load_labels(path):
    with open(path, 'r', encoding='utf-8') as f:
        return [l.strip() for l in f if l.strip()]


def preprocess_image(path, target_size=(224, 224)):
    img = Image.open(path).convert('RGB').resize(target_size)
    arr = np.array(img).astype('float32') / 255.0
    return arr[np.newaxis, ...]


def predict_with_model(model, image_path, labels, topk=5):
    # infer target size from model if available
    try:
        ishape = model.input_shape
    except Exception:
        ishape = None
    if ishape and len(ishape) >= 3 and all(isinstance(x, int) for x in ishape[1:3]):
        target = (ishape[1], ishape[2])
    else:
        target = (224, 224)
    x = preprocess_image(image_path, target_size=target)
    preds = model.predict(x)[0]
    idxs = np.argsort(preds)[-topk:][::-1]
    results = []
    for i in idxs:
        lab = labels[i] if i < len(labels) else str(i)
        results.append((int(i), float(preds[i]), lab))
    return results


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--image', '-i', required=True, help='Path to input image')
    p.add_argument('--model', '-m', default='birds_mobilenetv2.keras', help='Keras model file')
    p.add_argument('--labels', '-l', default='labels.txt', help='Labels file (one label per line)')
    p.add_argument('--topk', type=int, default=5, help='Show top-K predictions')
    args = p.parse_args()

    if not os.path.exists(args.image):
        print('Image not found:', args.image, file=sys.stderr)
        sys.exit(2)
    if not os.path.exists(args.model):
        print('Model not found:', args.model, file=sys.stderr)
        sys.exit(2)
    if not os.path.exists(args.labels):
        print('Labels file not found:', args.labels, file=sys.stderr)
        sys.exit(2)

    labels = load_labels(args.labels)

    # Load model early so we can validate label length vs model output size
    try:
        model = load_model(args.model)
    except Exception as e:
        print('Failed to load model:', e, file=sys.stderr)
        sys.exit(2)

    # Check label count vs model output
    try:
        out_dim = int(model.output_shape[-1])
    except Exception:
        out_dim = None

    if out_dim is not None and len(labels) != out_dim:
        print(f"Warning: labels file has {len(labels)} entries but model output dim is {out_dim}.", file=sys.stderr)
        print('If labels are numeric IDs, consider using a human-readable labels file matching training order.', file=sys.stderr)

    # Detect numeric-only labels (common when folder names are numeric IDs)
    if all(l.isdigit() for l in labels):
        print('Note: labels appear to be numeric IDs; model was likely trained with human-readable class names.', file=sys.stderr)
        print('Consider providing a labels file with species names or fetch labels from the original dataset.', file=sys.stderr)

    results = predict_with_model(model, args.image, labels, topk=args.topk)
    for rank, (idx, prob, label) in enumerate(results, start=1):
        print(f"{rank}. {label} (index={idx}) — {prob:.4f}")


if __name__ == '__main__':
    main()