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) }