diplom

TriView3D — Multi-View 3D Cell Reconstruction

Reconstruction of 3D cell morphology from 3 orthogonal projections (top, bottom, side).

Architecture

Base CNN (best_autoencoder.pt)

• Encoder2D: 2D ConvNet [3, 64, 64] → latent (256)
• Decoder3D: 3D ConvTranspose → [1, 64, 64, 64] with skip connections from lifted views
• Loss: BCE + Dice + Projection Consistency (BCE) + Boundary-aware BCE
• Training: AdamW + 5-epoch LR warmup + complexity-weighted sampling + augmentation

Detail Refiner (best_refiner.pt)

• Uncertainty-gated residual refinement on top of coarse CNN prediction
• Input: coarse logits + lifted views → gating * residual
• Skip connection version: compatible with skip-connected autoencoder

Conditional VAE (best_vae.pt)

• Encoder: 2D ConvNet → mu (256) + logvar (256)
• Decoder: 3D ConvTranspose with skip connections
• Inference: best-of-K sampling (K=8) by reprojection error
• Loss: Reconstruction + KL divergence

Classifier (best_classifier.pt)

• MLP on latent vector (256) + morphometric features (5: volume, sphericity, convexity, surface_area, compactness)
• Task: normal vs anomaly classification

Input Modes

• tri: 3 channels (top_proj, bottom_proj, side_proj) — primary mode
• quad: 4 channels (top_proj, bottom_proj, side_proj, front_proj)

Metrics (previous version, before skip connections + bug fixes)

• Dice: ~0.90
• IoU: ~0.82

Note: Models uploaded before the pipeline overhaul may have lower quality. Retrained models with skip connections, TTA, augmentation, and bug fixes are expected to achieve significantly higher Dice.

Usage

import torch
from src.autoencoder import TriViewAutoencoder
from src.reconstruction_utils import infer_in_channels_from_state_dict, infer_skip_channels_from_state_dict

state_dict = torch.load("best_autoencoder.pt", map_location="cpu")
if "model_state_dict" in state_dict:
    state_dict = state_dict["model_state_dict"]

in_channels = infer_in_channels_from_state_dict(state_dict)
skip_channels = infer_skip_channels_from_state_dict(state_dict)
latent_dim = int(state_dict["encoder.fc.1.weight"].shape[0])

model = TriViewAutoencoder(latent_dim=latent_dim, in_channels=in_channels, skip_channels=skip_channels)
model.load_state_dict(state_dict)
model.eval()

# Inference with TTA
import torch.nn.functional as F
x = ...  # [batch, 3, 64, 64] projections
pred = model(x)
pred_flipped = model(x.flip(-1)).flip(-1)
pred_avg = (pred + pred_flipped) / 2.0
volume = torch.sigmoid(pred_avg)  # [batch, 1, 64, 64, 64]

Pipeline Order

1. train_colab.ipynb → best_autoencoder.pt
2. train_refiner_colab.ipynb → best_refiner.pt (needs stage 1)
3. train_vae_colab.ipynb → best_vae.pt (independent)
4. train_classifier_colab.ipynb → best_classifier.pt (needs stage 1)