src/inference.py

"""Load and run the packaged model locally. No network calls are made."""
import json, h5py, numpy as np, torch
from pathlib import Path
from .model import InsectCNN
from .audio import logmel, audio_quality, resample, crop_offsets, read_audio, SR

# Labels are 'GenusSpecies' with no separator; longest prefix first so Pseudochorthippus != Chorthippus.
ORTHOPTERA=['Pseudochorthippus','Chorthippus','Gryllus','Nemobius','Oecanthus','Pholidoptera','Roeseliana','Tettigonia']
CICADIDAE=['Azanicada','Brevisiana','Kikihia','Myopsalta','Platypleura','Pycna']
# Chosen on the validation split for full_model (species ~77%, genus ~90%, group ~98% correct above it).
DEFAULT_THRESHOLD=0.6

def genus_of(species):
    return next((g for g in ORTHOPTERA+CICADIDAE if species.startswith(g)),species)

def group_of(species):
    return 'Orthoptera (crickets/grasshoppers)' if genus_of(species) in ORTHOPTERA else 'Cicadidae (cicadas)'

def _summed(probs, labels, keyf):
    out={}
    for i,p in enumerate(probs): k=keyf(labels[str(i)]); out[k]=out.get(k,0.)+float(p)
    return sorted(out.items(),key=lambda kv:-kv[1])

def identify(probs, labels, threshold=DEFAULT_THRESHOLD):
    """Most specific taxonomic level whose summed probability reaches the threshold, else 'unknown'."""
    top=int(np.argmax(probs))
    levels=[('species',[(labels[str(top)],float(probs[top]))]),
            ('genus',_summed(probs,labels,genus_of)),('group',_summed(probs,labels,group_of))]
    for level,ranked in levels:
        if ranked[0][1]>=threshold: return {'level':level,'name':ranked[0][0],'confidence':ranked[0][1]}
    return {'level':'unknown','name':'Uncertain - possibly a species outside the 32 trained classes','confidence':levels[-1][1][0][1]}

def load_model(model_dir, device='cpu'):
    model_dir=Path(model_dir)
    cfg=json.loads((model_dir/'config.json').read_text())
    labels=json.loads((model_dir/'labels.json').read_text())
    net=InsectCNN(cfg['num_classes'])
    h5_path=model_dir/'model.h5'
    with h5py.File(h5_path,'r') as h:
        state={}
        for key in h['state_dict']:
            arr=np.asarray(h['state_dict'][key])
            state[key]=torch.from_numpy(arr.copy())
    net.load_state_dict(state)
    net.to(device).eval()
    return net, labels, cfg

def predict_wav(path, model_dir, device='cpu'):
    net,labels,cfg=load_model(model_dir,device)
    # Integer PCM is scaled to [-1,1] before mixing to mono so the quality check sees true levels.
    sr,y=read_audio(path)
    ok,reasons=audio_quality(y)
    # eval_crops>1 averages over evenly spaced 2 s windows; older bundles use the center crop.
    y16=resample(y,sr)
    offs=crop_offsets(len(y16),int(cfg.get('eval_crops',1)))
    x=torch.from_numpy(np.stack([logmel(y16,SR,o) for o in offs])).to(device)
    with torch.no_grad():
        probs=torch.softmax(net(x),dim=1).mean(dim=0).cpu().numpy()
    order=np.argsort(probs)[::-1][:5]
    thr=float(cfg.get('confidence_threshold',DEFAULT_THRESHOLD))
    return {'identification':identify(probs,labels,thr),
            'genus_scores':[{'genus':g,'score':v} for g,v in _summed(probs,labels,genus_of)[:3]],
            'group_scores':[{'group':g,'score':v} for g,v in _summed(probs,labels,group_of)],
            'predicted_species':labels[str(int(order[0]))],
            'confidence':float(probs[order[0]]),
            'top_k':[{'species':labels[str(int(i))],'score':float(probs[i])} for i in order],
            'quality_pass':bool(ok),'quality_reasons':reasons,
            'model_version':cfg['model_version'],
            'experimental':True,
            'warning':'Research insect-species classifier only. Not trained for grain-pest detection or grain safety.'}