| """ |
| ================================================================= |
| TubuleSegmentation v1.0.0 ================================================================= |
| - EfficientNet-B4 encoder, depth=5, ImageNet weights |
| - UnetDecoder (256,128,64,32,16), n_blocks=5, attention='scse' |
| - Macenko stain aug (p=0.5), gradient checkpointing encoder+decoders |
| - Dice+CE + clDice(0.04) + SCNP(0.16) + Containment(0.24) |
| - Warmup topo lineal epochs 20-40 |
| - batch=4, LR=3e-4, AdamW, CosineAnnealingWarmRestarts |
| |
| Output: best_model_v1.0.pt |
| ================================================================= |
| """ |
| import os, time, numpy as np, cv2, torch, torch.nn as nn |
| import torch.nn.functional as F |
| import torch.utils.checkpoint as cp |
| import segmentation_models_pytorch as smp |
| from segmentation_models_pytorch.decoders.unet.decoder import UnetDecoder |
| from torch.utils.data import Dataset, DataLoader |
| import albumentations as A |
| from pathlib import Path |
|
|
| |
|
|
| TORCHSTAIN_OK = False |
| TorchMacenkoAugmentor = None |
| try: |
| from torchstain.torch.augmentors.macenko import TorchMacenkoAugmentor |
| TORCHSTAIN_OK = True |
| except ImportError: |
| try: |
| from torchstain.augmentors.he_augmentor import HEAugmentor as TorchMacenkoAugmentor |
| TORCHSTAIN_OK = True |
| except ImportError: |
| print(" [!] torchstain no encontrado — pip install torchstain") |
| print(" [!] Entrenando sin stain augmentation (igual que v21)") |
|
|
| |
| |
| |
| def scnp(logits, labels_onehot, w=3, kappa=1e6): |
| fg_mask = labels_onehot.float(); bg_mask = 1.0 - fg_mask |
| fg_logits = logits * fg_mask + kappa * bg_mask |
| min_pooled = -F.max_pool2d(-fg_logits, kernel_size=w, stride=1, padding=w//2) |
| bg_logits = logits * bg_mask - kappa * fg_mask |
| max_pooled = F.max_pool2d(bg_logits, kernel_size=w, stride=1, padding=w//2) |
| return torch.where(labels_onehot == 1, min_pooled, max_pooled) |
|
|
| def soft_erode(img): |
| p1 = -F.max_pool2d(-img, (3,1), (1,1), (1,0)) |
| p2 = -F.max_pool2d(-img, (1,3), (1,1), (0,1)) |
| return torch.min(p1, p2) |
|
|
| def soft_dilate(img): |
| return F.max_pool2d(img, (3,3), (1,1), (1,1)) |
|
|
| def soft_open(img): |
| return soft_dilate(soft_erode(img)) |
|
|
| def soft_skel(img, iters=3): |
| img1 = soft_open(img); skel = F.relu(img - img1) |
| for _ in range(iters): |
| img = soft_erode(img); img1 = soft_open(img) |
| delta = F.relu(img - img1); skel = skel + F.relu(delta - skel * delta) |
| return skel |
|
|
| def cldice_loss(pred, target, iters=3, smooth=1e-7): |
| skel_pred = soft_skel(pred, iters); skel_target = soft_skel(target, iters) |
| tprec = (skel_pred * target).sum() / (skel_pred.sum() + smooth) |
| tsens = (skel_target * pred).sum() / (skel_target.sum() + smooth) |
| return 1.0 - 2.0 * tprec * tsens / (tprec + tsens + smooth) |
|
|
| def containment_loss(pred_lumen, pred_epi, smooth=1e-5): |
| tissue = torch.clamp(pred_epi + pred_lumen, 0, 1) |
| violation = pred_lumen * (1.0 - tissue.detach()) |
| return violation.sum() / (pred_lumen.sum() + smooth) |
|
|
| |
| |
| |
| class TubuleSegModel(nn.Module): |
| def __init__(self, use_checkpoint=True): |
| super().__init__() |
| self.use_checkpoint = use_checkpoint |
| self.encoder = smp.encoders.get_encoder( |
| 'timm-efficientnet-b4', in_channels=3, depth=5, weights='imagenet' |
| ) |
| 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='scse' |
| ) |
| 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='scse' |
| ) |
| self.border_head = nn.Conv2d(16, 2, kernel_size=1) |
|
|
| def _seg_branch(self, *features): |
| return self.seg_head(self.seg_decoder(list(features))) |
|
|
| def _border_branch(self, *features): |
| return self.border_head(self.border_decoder(list(features))) |
|
|
| def _encode(self, x): |
| return tuple(self.encoder(x)) |
|
|
| def forward(self, x): |
| if self.use_checkpoint and self.training: |
| features = list(cp.checkpoint(self._encode, x, use_reentrant=False)) |
| seg = cp.checkpoint(self._seg_branch, *features, use_reentrant=False) |
| border = cp.checkpoint(self._border_branch, *features, use_reentrant=False) |
| return seg, border |
| features = self.encoder(x) |
| return ( |
| self.seg_head(self.seg_decoder(features)), |
| self.border_head(self.border_decoder(features)) |
| ) |
|
|
| |
| |
| |
| def multiclass_dice_loss(pred_logits, target_long, smooth=1e-5): |
| C = pred_logits.shape[1]; soft = pred_logits.softmax(dim=1) |
| oh = F.one_hot(target_long, C).permute(0,3,1,2).float() |
| num = 2.0*(soft*oh).sum(dim=(0,2,3))+smooth |
| den = soft.sum(dim=(0,2,3))+oh.sum(dim=(0,2,3))+smooth |
| return (1.0-(num/den)).mean() |
|
|
| def seg_loss_base(logits, target): |
| return multiclass_dice_loss(logits, target) + F.cross_entropy(logits, target) |
|
|
| |
| |
| |
| def extract_training_data(data_repo, masks_dir): |
| orig_dir=data_repo/'tubules_original';conteo_dir=data_repo/'tubules_area_ok' |
| images=[];masks=[];borders_outer=[];borders_lumen=[] |
| border_pos=np.zeros(2);border_total=0 |
| n_loaded=0;n_missing=0;_printed_unique=[False] |
| def process(orig_path,conteo_path): |
| nonlocal border_pos,border_total,n_loaded,n_missing |
| orig=cv2.imread(str(orig_path));conteo=cv2.imread(str(conteo_path)) |
| if orig is None or conteo is None:return |
| oh,ow=orig.shape[:2];ch,cw=conteo.shape[:2] |
| r=conteo[:,:,2].astype(float);g=conteo[:,:,1].astype(float);b=conteo[:,:,0].astype(float) |
| red=((r>150)&(g<100)&(b<100)&((r-g)>80)).astype(np.uint8) |
| if ch>oh:red=red[50:ch-50,50:cw-50] |
| red=cv2.resize(red,(ow,oh),interpolation=cv2.INTER_NEAREST) |
| |
| bid=Path(orig_path).stem |
| mask_path=masks_dir/f"{bid}.png" |
| if not mask_path.exists(): |
| n_missing+=1;return |
| seg_mask=cv2.imread(str(mask_path),cv2.IMREAD_GRAYSCALE) |
| if seg_mask is None: |
| n_missing+=1;return |
| if seg_mask.shape[:2]!=(oh,ow): |
| seg_mask=cv2.resize(seg_mask,(ow,oh),interpolation=cv2.INTER_NEAREST) |
| seg_mask=seg_mask.astype(np.uint8) |
| if not _printed_unique[0]: |
| print(f" [v23] {bid}.png -> labels {np.unique(seg_mask).tolist()} (esperado [0,1,2])") |
| _printed_unique[0]=True |
| if (seg_mask==1).sum()<oh*ow*0.1:return |
| n_loaded+=1 |
| barrier=cv2.dilate(red,cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)),iterations=1) |
| flood=barrier.copy();fm=np.zeros((oh+2,ow+2),dtype=np.uint8) |
| for y in range(0,oh,5): |
| if flood[y,0]==0:cv2.floodFill(flood,fm,(0,y),128) |
| if flood[y,ow-1]==0:cv2.floodFill(flood,fm,(ow-1,y),128) |
| for x in range(0,ow,5): |
| if flood[0,x]==0:cv2.floodFill(flood,fm,(x,0),128) |
| if flood[oh-1,x]==0:cv2.floodFill(flood,fm,(x,oh-1),128) |
| bg=(flood==128).astype(np.uint8) |
| bg_dil=cv2.dilate(bg,cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(7,7))) |
| ob=(red&bg_dil).astype(np.float32);lb=(red&(1-bg_dil)).astype(np.float32) |
| k5=cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5)) |
| ob=cv2.dilate(ob,k5,iterations=1);lb=cv2.dilate(lb,k5,iterations=1) |
| images.append(cv2.resize(orig,(512,512),interpolation=cv2.INTER_LINEAR)) |
| masks.append(cv2.resize(seg_mask,(512,512),interpolation=cv2.INTER_NEAREST)) |
| borders_outer.append(cv2.resize(ob,(512,512),interpolation=cv2.INTER_NEAREST)) |
| borders_lumen.append(cv2.resize(lb,(512,512),interpolation=cv2.INTER_NEAREST)) |
| border_pos[0]+=borders_outer[-1].sum();border_pos[1]+=borders_lumen[-1].sum() |
| border_total+=512*512 |
| for ap in sorted(conteo_dir.glob('*_conteo.png')): |
| bid=get_base_id(ap.name) |
| if bid.split('_')[-1]=='396':continue |
| op=orig_dir/f"{bid}.tif" |
| if not op.exists():continue |
| process(op,ap) |
| for ext_name in ['tubules_extra','tubules_extra2','tubules_extra4']: |
| ext_dir=data_repo/ext_name |
| if not ext_dir.exists():continue |
| for tp in sorted(ext_dir.glob('*.tif')): |
| bid=tp.stem;cp_path=ext_dir/f"{bid}_conteo.png" |
| if not cp_path.exists():continue |
| process(tp,cp_path) |
| pw=np.clip((border_total-border_pos)/(border_pos+1e-6),1.0,15.0) if border_total>0 else np.array([10.,10.]) |
| print(f" [v23] mascaras cargadas: {n_loaded} | sin mascara (omitidas): {n_missing}") |
| return images,masks,borders_outer,borders_lumen,pw |
|
|
| def get_base_id(f): |
| s=os.path.splitext(f)[0] |
| for x in['_medida','_conteo','_area']: |
| if s.endswith(x):s=s[:-len(x)] |
| return s |
|
|
| |
| |
| |
| def make_stain_augmentor(): |
| |
| if not TORCHSTAIN_OK: |
| return None |
| try: |
| aug = TorchMacenkoAugmentor(sigma1=0.2, sigma2=0.2) |
| return aug |
| except Exception as e: |
| print(f" [!] Macenko augmentor init fallo: {e}") |
| return None |
|
|
| def apply_stain_aug(img_rgb_uint8, augmentor): |
| |
| if augmentor is None: |
| return img_rgb_uint8 |
| try: |
| |
| img_t = torch.from_numpy(img_rgb_uint8).permute(2, 0, 1) |
| img_aug, _, _ = augmentor(img_t) |
| |
| img_out = img_aug.permute(1, 2, 0).numpy().astype(np.uint8) |
| return img_out |
| except Exception: |
| return img_rgb_uint8 |
|
|
| |
| |
| |
| class TubuleDataset(Dataset): |
| def __init__(self, images, masks, borders_o, borders_l, |
| augment=False, stain_augmentor=None): |
| self.images = images |
| self.masks = masks |
| self.borders_o = borders_o |
| self.borders_l = borders_l |
| self.augment = augment |
| self.stain_augmentor = stain_augmentor |
|
|
| self.tf = A.Compose([ |
| A.HorizontalFlip(p=0.5), |
| A.VerticalFlip(p=0.5), |
| A.RandomRotate90(p=0.5), |
| A.ElasticTransform(alpha=120, sigma=120*0.05, p=0.3), |
| A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05, p=0.8), |
| A.GaussNoise(p=0.3), |
| ]) if augment else None |
|
|
| def __len__(self): |
| return len(self.images) |
|
|
| def __getitem__(self, idx): |
| img = cv2.cvtColor(self.images[idx], cv2.COLOR_BGR2RGB) |
| mask = self.masks[idx].copy() |
| bo = self.borders_o[idx].copy() |
| bl = self.borders_l[idx].copy() |
|
|
| |
| |
| if self.stain_augmentor is not None and np.random.random() < 0.5: |
| img = apply_stain_aug(img, self.stain_augmentor) |
|
|
| if self.tf: |
| bou = (bo * 255).astype(np.uint8) |
| blu = (bl * 255).astype(np.uint8) |
| t = self.tf(image=img, masks=[mask, bou, blu]) |
| img = t['image']; mask = t['masks'][0] |
| bo = t['masks'][1].astype(np.float32) / 255.0 |
| bl = t['masks'][2].astype(np.float32) / 255.0 |
|
|
| img_norm = (img.astype(np.float32)/255.0 - [0.485,0.456,0.406]) / [0.229,0.224,0.225] |
| return ( |
| torch.from_numpy(img_norm.transpose(2,0,1)).float(), |
| torch.from_numpy(mask.copy()).long(), |
| torch.from_numpy(np.stack([bo, bl], axis=0)).float() |
| ) |
|
|
| |
| |
| |
| def main(): |
| base_dir = Path(r"D:\Lu\AI\Tubules"); os.chdir(str(base_dir)) |
| print("="*60) |
| print(" TubuleSegmentation v23") |
| print(" EfficientNet-B4 @ 512x512 | batch=4 | SCSE attention") |
| print(" Dice+CE + clDice(0.04) + SCNP(0.16) + Containment(0.24)") |
| print(" Macenko stain aug (p=0.5) + checkpointing encoder+decoders") |
| print(f" torchstain: {'v1.3.0 OK' if TORCHSTAIN_OK else 'NO (pip install torchstain)'}") |
| print("="*60) |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| print(f"\n Device: {device}") |
| if device.type == 'cuda': |
| print(f" GPU: {torch.cuda.get_device_name(0)}") |
| torch.backends.cudnn.benchmark = True |
|
|
| 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)) |
|
|
| masks_dir = base_dir / 'masks_v2' / 'masks' |
| |
| print("\n Loading images + masks_v2...") |
| images, masks, borders_o, borders_l, border_pw = extract_training_data(data_repo, masks_dir) |
| print(f" Total images: {len(images)}") |
| print(f" Border pos_weight: outer={border_pw[0]:.1f}, lumen={border_pw[1]:.1f}") |
|
|
| np.random.seed(42) |
| idx = np.random.permutation(len(images)) |
| sp = int(len(idx) * 0.85) |
| train_idx = idx[:sp]; val_idx = idx[sp:] |
| print(f" Train: {len(train_idx)}, Val: {len(val_idx)}") |
|
|
| stain_aug = make_stain_augmentor() |
| if stain_aug is not None: |
| print(" Macenko stain augmentation: YES (sigma=0.2, p=0.5 por imagen)") |
| else: |
| print(" Macenko stain augmentation: NO") |
|
|
| train_ds = TubuleDataset( |
| [images[i] for i in train_idx], [masks[i] for i in train_idx], |
| [borders_o[i] for i in train_idx], [borders_l[i] for i in train_idx], |
| augment=True, stain_augmentor=stain_aug |
| ) |
| val_ds = TubuleDataset( |
| [images[i] for i in val_idx], [masks[i] for i in val_idx], |
| [borders_o[i] for i in val_idx], [borders_l[i] for i in val_idx], |
| augment=False, stain_augmentor=None |
| ) |
|
|
| BATCH = 4 |
| train_dl = DataLoader(train_ds, batch_size=BATCH, shuffle=True, num_workers=0, pin_memory=True) |
| val_dl = DataLoader(val_ds, batch_size=BATCH, shuffle=False, num_workers=0, pin_memory=True) |
|
|
| model = TubuleSegModel(use_checkpoint=True); model.to(device) |
| n_params = sum(p.numel() for p in model.parameters()) / 1e6 |
| print(f"\n Model: {n_params:.1f}M params (SCSE attention en decoders)") |
|
|
| |
| model.train() |
| if device.type == 'cuda': torch.cuda.reset_peak_memory_stats() |
| _x = torch.randn(2, 3, 512, 512, device=device) |
| _seg, _bor = model(_x) |
| assert _seg.shape == (2,3,512,512) and _bor.shape == (2,2,512,512) |
| (_seg.mean() + _bor.mean()).backward() |
| assert any(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()), "grad no finito" |
| model.zero_grad(set_to_none=True) |
| if device.type == 'cuda': |
| print(f" Smoke test OK | VRAM_peak smoke={torch.cuda.max_memory_allocated()/1024**2:.0f}MB (batch=2)") |
| torch.cuda.reset_peak_memory_stats() |
| else: |
| print(" Smoke test OK") |
|
|
| W_CLDICE = 0.04; W_SCNP = 0.16; W_CONTAIN = 0.24 |
| TOPO_START = 20; TOPO_END = 40 |
| EPOCHS = 200; PATIENCE = 80; LR = 3e-4 |
|
|
| opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4) |
| sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=20, T_mult=2, eta_min=1e-6) |
| border_pw_t = torch.tensor(border_pw, dtype=torch.float32).to(device) |
|
|
| best = 0; pat = 0 |
| model_path = base_dir / 'best_model_v23.pt' |
|
|
| print(f"\n Epochs: {EPOCHS}, Patience: {PATIENCE}, LR: {LR}") |
| print(f" Topo: clDice={W_CLDICE}, SCNP={W_SCNP}, Contain={W_CONTAIN}") |
| print(f" Warmup: epochs {TOPO_START}-{TOPO_END} (linear 0->1) | Batch: {BATCH}") |
| print(f"{'='*60}") |
| t0 = time.time() |
|
|
| for ep in range(EPOCHS): |
| if ep < TOPO_START: topo_scale = 0.0 |
| elif ep < TOPO_END: topo_scale = (ep - TOPO_START) / (TOPO_END - TOPO_START) |
| else: topo_scale = 1.0 |
|
|
| if device.type == 'cuda': torch.cuda.reset_peak_memory_stats() |
| model.train(); tl = 0; tb = 0 |
|
|
| for imgs, seg_masks, borders in train_dl: |
| imgs = imgs.to(device) |
| seg_masks = seg_masks.to(device) |
| borders = borders.to(device) |
| opt.zero_grad() |
|
|
| seg_logits, border_logits = model(imgs) |
| l_base = seg_loss_base(seg_logits, seg_masks) |
| l_border = F.binary_cross_entropy_with_logits( |
| border_logits.float(), borders.float(), |
| pos_weight=border_pw_t.view(1,2,1,1)) |
| loss = l_base + 0.2 * l_border |
|
|
| if topo_scale > 0: |
| seg_probs = seg_logits.softmax(dim=1) |
| epi_pred = seg_probs[:,1:2,:,:] |
| epi_target = (seg_masks==1).float().unsqueeze(1) |
| l_cldice = cldice_loss(epi_pred, epi_target, iters=3) |
| labels_oh = F.one_hot(seg_masks, 3).permute(0,3,1,2).float() |
| z_tilde = scnp(seg_logits, labels_oh, w=3) |
| l_scnp = F.cross_entropy(z_tilde, seg_masks) |
| l_contain = containment_loss(seg_probs[:,2,:,:], seg_probs[:,1,:,:]) |
| loss = loss + topo_scale * (W_CLDICE*l_cldice + W_SCNP*l_scnp + W_CONTAIN*l_contain) |
|
|
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) |
| opt.step() |
| tl += loss.item(); tb += 1 |
|
|
| sched.step() |
|
|
| model.eval(); ap = []; at = [] |
| with torch.no_grad(): |
| for imgs, seg_masks, _ in val_dl: |
| imgs = imgs.to(device) |
| seg_logits, _ = model(imgs) |
| ap.append(seg_logits.argmax(1).cpu()) |
| at.append(seg_masks) |
|
|
| preds = torch.cat(ap); tgts = torch.cat(at) |
| ious = [(((preds==c)&(tgts==c)).sum().float()+1e-6) / |
| (((preds==c)|(tgts==c)).sum().float()+1e-6) for c in range(3)] |
| miou = np.mean([x.item() for x in ious]) |
|
|
| peak_mb = (torch.cuda.max_memory_allocated()/1024**2) if device.type=='cuda' else 0 |
| if (ep+1) % 5 == 0 or miou > best: |
| print(f" Ep {ep+1:3d}/{EPOCHS} | loss={tl/tb:.4f} | " |
| f"mIoU={miou:.4f} epi={ious[1]:.3f} lum={ious[2]:.3f} | " |
| f"{(time.time()-t0)/60:.1f}m | topo={topo_scale:.2f} | " |
| f"VRAM_peak={peak_mb:.0f}MB") |
| if miou > best: |
| best = miou; pat = 0 |
| torch.save({ |
| 'model_state_dict': model.state_dict(), |
| 'config': { |
| 'encoder': 'efficientnet-b4', |
| 'architecture': 'UNet_FullImage_ShapeConstrained_SCSE', |
| 'img_size': 512, |
| 'losses': 'Dice+CE+clDice+SCNP+Containment', |
| 'attention': 'scse', |
| 'stain_aug': 'macenko_sigma0.2' if stain_aug else 'none', |
| 'gradient_checkpointing': 'encoder_full+decoders', |
| 'topo_weights': {'cldice': W_CLDICE, 'scnp': W_SCNP, 'containment': W_CONTAIN}, |
| 'topo_warmup': 'linear_20_40', |
| 'from_scratch': True, |
| 'clean_dataset': True, |
| 'mask_source': 'masks_v2/masks (regenerate, +107 con --no-is-clean)', |
| 'version': 'v23' |
| }, |
| 'best_val_iou': best |
| }, str(model_path)) |
| print(f" * MEJOR: {best:.4f}") |
| else: |
| pat += 1 |
| if pat >= PATIENCE: |
| print(f" Early stopping at epoch {ep+1} (best={best:.4f})") |
| break |
|
|
| print(f"\n{'='*60}\n COMPLETE in {(time.time()-t0)/60:.1f} min — mIoU: {best:.4f}\n{'='*60}") |
| try: |
| from huggingface_hub import HfApi; api = HfApi() |
| api.upload_file( |
| path_or_fileobj=str(model_path), |
| path_in_repo='best_model_v23.pt', |
| repo_id='LuGot16/seminiferous-tubule-segmentation', |
| repo_type='model') |
| print(" Upload to HF!") |
| except Exception as e: |
| print(f" Upload: {e}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
