inference_v1.0.0.py

"""
=================================================================
INFERENCE v1.0.0 — Full-image + SCSE attention + TTA
=================================================================
"""

import os, cv2, torch, numpy as np, json, time, csv
import torch.nn as nn
import torch.nn.functional as F
import segmentation_models_pytorch as smp
from segmentation_models_pytorch.decoders.unet.decoder import UnetDecoder
from pathlib import Path

# ============================================================
# MODEL v1.0.0
# ============================================================
class TubuleSegModel(nn.Module):
    def __init__(self, attention_type='scse'):
        super().__init__()
        self.encoder = smp.encoders.get_encoder(
            'timm-efficientnet-b4', in_channels=3, depth=5, weights=None
        )
        ec = self.encoder.out_channels
        self.seg_decoder = UnetDecoder(
            encoder_channels=ec, decoder_channels=(256,128,64,32,16),
            n_blocks=5, use_norm='batchnorm', attention_type=attention_type
        )
        self.seg_head = nn.Conv2d(16, 3, kernel_size=1)
        self.border_decoder = UnetDecoder(
            encoder_channels=ec, decoder_channels=(256,128,64,32,16),
            n_blocks=5, use_norm='batchnorm', attention_type=attention_type
        )
        self.border_head = nn.Conv2d(16, 2, kernel_size=1)

    def forward(self, x):
        features = self.encoder(x)
        return (
            self.seg_head(self.seg_decoder(features)),
            self.border_head(self.border_decoder(features))
        )


# ============================================================
# TTA — Test-Time Augmentation
# ============================================================
def predict_with_tta(model, img_t, device):
    """
    8 pasadas: 4 rotaciones (0, 90, 180, 270) x 2 flips (original, horizontal).
    Promedia probabilidades softmax en el espacio original.
    img_t: [1, 3, H, W] tensor normalizado en device.
    Retorna: [3, H, W] numpy array de probabilidades promediadas.
    """
    preds = []
    with torch.no_grad():
        for flip in [False, True]:
            x = img_t.flip(-1) if flip else img_t
            for k in range(4):  # rotaciones 0, 90, 180, 270
                x_rot = torch.rot90(x, k, dims=[-2, -1])
                seg_logits, _ = model(x_rot)
                prob = F.softmax(seg_logits, dim=1)  # [1, 3, H, W]
                # Revertir la rotacion
                prob = torch.rot90(prob, -k, dims=[-2, -1])
                # Revertir el flip
                if flip:
                    prob = prob.flip(-1)
                preds.append(prob.cpu().numpy()[0])  # [3, H, W]

    return np.mean(preds, axis=0)  # [3, H, W]


# ============================================================
# POST-PROCESSING (identico a inference_v22.py)
# ============================================================
def postprocess(pred, h, w):
    epi_mask = (pred == 1).astype(np.uint8)
    n, labels = cv2.connectedComponents(epi_mask)
    if n > 1:
        areas = [(lid, (labels==lid).sum()) for lid in range(1, n)]
        areas.sort(key=lambda x: x[1], reverse=True)
        epi_clean = (labels == areas[0][0]).astype(np.uint8)
    else:
        epi_clean = epi_mask

    k_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
    epi_closed = cv2.morphologyEx(epi_clean, cv2.MORPH_CLOSE, k_close)

    flood = epi_closed.copy()
    fm = np.zeros((h+2, w+2), dtype=np.uint8)
    for y in range(0, h, 3):
        if flood[y, 0] == 0:   cv2.floodFill(flood, fm, (0,   y), 2)
        if flood[y, w-1] == 0: cv2.floodFill(flood, fm, (w-1, y), 2)
    for x in range(0, w, 3):
        if flood[0, x] == 0:   cv2.floodFill(flood, fm, (x,   0), 2)
        if flood[h-1, x] == 0: cv2.floodFill(flood, fm, (x, h-1), 2)
    holes = (flood == 0).astype(np.uint8)
    epi_filled = epi_closed | holes

    final = np.zeros((h, w), dtype=np.uint8)
    final[epi_filled > 0] = 1

    lum_mask = (pred == 2).astype(np.uint8)
    lum_inside = lum_mask & epi_filled
    n, labels = cv2.connectedComponents(lum_inside)
    if n > 1:
        lum_clean = np.zeros_like(lum_inside)
        for lid in range(1, n):
            comp = (labels == lid).astype(np.uint8)
            if comp.sum() < 300: continue
            dilated = cv2.dilate(comp, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5)))
            border_px = dilated & (~epi_filled.astype(bool)).astype(np.uint8)
            if border_px.sum() < comp.sum() * 0.05:
                lum_clean[labels == lid] = 1
        final[lum_clean > 0] = 2

    return final


# ============================================================
# MAIN
# ============================================================
def main():
    base_dir = Path(r"D:\Lu\AI\Tubules"); os.chdir(str(base_dir))
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")

    model_file = 'best_model_segmentation_v1.0.0.pt'
    attention  = 'scse'

    model = TubuleSegModel(attention_type=attention)
    ckpt  = torch.load(str(base_dir / model_file), map_location='cpu', weights_only=False)
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device); model.eval()
    
    from huggingface_hub import snapshot_download
    data_repo = base_dir / 'data_repo'
    if not data_repo.exists():
        snapshot_download('LuGot16/tubules', repo_type='dataset', local_dir=str(data_repo))

    test_dir = data_repo / 'area_test'
    out_dir  = base_dir / f'inference_results_{model_file.replace("best_model_","").replace(".pt","")}'
    out_dir.mkdir(exist_ok=True)

    SCALE   = 0.32  # um/pixel
    results = []
    files   = sorted(test_dir.glob('*.tif'))
    t0 = time.time()

    for i, fp in enumerate(files):
        img = cv2.imread(str(fp))
        if img is None: continue
        h, w = img.shape[:2]
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        img_resized = cv2.resize(img_rgb, (512, 512), interpolation=cv2.INTER_LINEAR)
        img_norm = (img_resized.astype(np.float32)/255.0 - [0.485,0.456,0.406]) / [0.229,0.224,0.225]
        img_t = torch.from_numpy(img_norm.transpose(2,0,1)).float().unsqueeze(0).to(device)

        # TTA
        seg_probs = predict_with_tta(model, img_t, device)  # [3, 512, 512]

        pred_512 = seg_probs.argmax(0).astype(np.uint8)
        pred     = cv2.resize(pred_512, (w, h), interpolation=cv2.INTER_NEAREST)
        pred     = postprocess(pred, h, w)

        total   = h * w
        epi_pct = (pred==1).sum() / total * 100
        lum_pct = (pred==2).sum() / total * 100
        tub_pct = epi_pct + lum_pct
        epi_um2 = (pred==1).sum() * SCALE * SCALE
        lum_um2 = (pred==2).sum() * SCALE * SCALE
        tub_um2 = epi_um2 + lum_um2
        status  = "OK" if tub_pct >= 30 and lum_pct > 1 else "CHECK"

        results.append({
            'image':          fp.name,
            'tubule_pct':     round(tub_pct, 1),
            'epithelium_pct': round(epi_pct, 1),
            'lumen_pct':      round(lum_pct, 1),
            'tubule_um2':     round(tub_um2, 1),
            'epithelium_um2': round(epi_um2, 1),
            'lumen_um2':      round(lum_um2, 1),
            'lumen_epi_ratio': round(lum_um2 / (epi_um2 + 1e-6), 4),
            'status':         status
        })

        # Overlay
        overlay = img.copy()
        overlay[pred==1] = (overlay[pred==1]*0.5 + np.array([0,180,0])*0.5).astype(np.uint8)
        overlay[pred==2] = (overlay[pred==2]*0.5 + np.array([255,100,0])*0.5).astype(np.uint8)

        tubule_mask = (pred >= 1).astype(np.uint8)
        contours_outer, _ = cv2.findContours(tubule_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(overlay, contours_outer, -1, (0, 255, 255), 2)

        lum_mask_vis = (pred == 2).astype(np.uint8)
        contours_lumen, _ = cv2.findContours(lum_mask_vis, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(overlay, contours_lumen, -1, (0, 0, 255), 2)

        cv2.imwrite(str(out_dir / f"{fp.stem}_overlay.png"), overlay)
        cv2.imwrite(str(out_dir / f"{fp.stem}_mask.png"), pred * 127)

        if (i+1) % 10 == 0 or status == "CHECK":
            print(f"  [{i+1}/{len(files)}] {fp.name}: "
                  f"tub={tub_pct:.1f}% epi={epi_pct:.1f}% lum={lum_pct:.1f}% [{status}]")

    ok = sum(1 for r in results if r['status'] == 'OK')
    tub_vals = [r['tubule_pct']     for r in results]
    epi_vals = [r['epithelium_pct'] for r in results]
    lum_vals = [r['lumen_pct']      for r in results]

    print(f"\n{'='*60}")
    print(f"  RESULTS {model_file} + TTA — {ok}/{len(results)} OK")
    print(f"  Tubule:     {np.mean(tub_vals):.1f} +/- {np.std(tub_vals):.1f}%")
    print(f"  Epithelium:   {np.mean(epi_vals):.1f} +/- {np.std(epi_vals):.1f}%")
    print(f"  Lumen:      {np.mean(lum_vals):.1f} +/- {np.std(lum_vals):.1f}%")
    print(f"  Time:     {(time.time()-t0)/60:.1f} min")
    print(f"{'='*60}")

    with open(out_dir / 'results.csv', 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=results[0].keys())
        w.writeheader(); w.writerows(results)
    with open(out_dir / 'results.json', 'w') as f:
        json.dump({
            'model': model_file, 'tta': True, 'n_tta': 8,
            'mIoU': float(ckpt['best_val_iou']),
            'n_ok': ok, 'n_total': len(results),
            'mean_tubule_pct': round(np.mean(tub_vals), 1),
            'mean_lumen_pct':  round(np.mean(lum_vals), 1),
            'results': results
        }, f, indent=2)
   


if __name__ == '__main__':
    main()