decision.py

import torch
import torch.nn.functional as F

MOLD_HIGH_CONF = 0.80
MOLD_LOW_CONF  = 0.50
BIO_CONF       = 0.60

def final_decision(model, img_tensor, mold_idx=4):
    with torch.no_grad():
        out = model(img_tensor.unsqueeze(0))
        class_probs = F.softmax(out["class"], dim=1)[0]
        bio_probs = F.softmax(out["bio"], dim=1)[0]

    mold_p = class_probs[mold_idx].item()
    bio_p  = bio_probs[1].item()

    if mold_p >= MOLD_HIGH_CONF:
        decision = "Mold"
    elif mold_p >= MOLD_LOW_CONF and bio_p >= BIO_CONF:
        decision = "Possible Mold"
    else:
        decision = "Not Mold"

    return {
        "decision": decision,
        "mold_probability": round(mold_p, 3),
        "biological_probability": round(bio_p, 3)
    }