mold-detection-api / advanced_decision.py
AdarshRajDS
Add mold detection FastAPI backend v2
e2f2323
import torch
import torch.nn.functional as F
import numpy as np
MOLD_HIGH = 0.85
MOLD_LOW = 0.50
BIO_TH = 0.70
UNCERT_TH = 0.15
PATCH_TH = 0.50
DINO_TH = 0.75
def enable_dropout(model):
for m in model.modules():
if m.__class__.__name__.startswith("Dropout"):
m.train()
def mc_uncertainty(model, img_tensor, mold_idx, T=15):
enable_dropout(model)
probs = []
for _ in range(T):
out = model(img_tensor.unsqueeze(0))
cp = F.softmax(out["class"],1)[0]
probs.append(cp[mold_idx].item())
model.eval()
return float(np.mean(probs)), float(np.std(probs))
def patch_consistency(model, image, transform, mold_idx, device, stride=112):
w,h = image.size
votes = []
for y in range(0, h-224+1, stride):
for x in range(0, w-224+1, stride):
patch = image.crop((x,y,x+224,y+224))
pt = transform(patch).unsqueeze(0).to(device)
with torch.no_grad():
out = model(pt)
cp = F.softmax(out["class"],1)[0]
votes.append(cp[mold_idx].item())
return float(np.mean(np.array(votes) > 0.7)) if votes else 0.0
def final_decision_v2(
mold_p, bio_p, uncert, patch_ratio, dino_sim
):
if (
mold_p > MOLD_HIGH and
bio_p > BIO_TH and
uncert < UNCERT_TH and
patch_ratio > PATCH_TH and
dino_sim > DINO_TH
):
return "Mold"
if mold_p > MOLD_LOW and bio_p > BIO_TH:
return "Possible Mold"
return "Not Mold"