add full model files

full_model/rnn_dataset.py

@@ -0,0 +1,178 @@
+import os
+import json
+import pydicom
+import numpy as np
+import torch
+
+from typing import Callable, Optional, Tuple
+
+from torch import Tensor
+from torch.utils.data import Dataset
+from sklearn.preprocessing import RobustScaler
+
+DTYPE = torch.float16
+
+
+class SyntaxDataset(Dataset):
+    def __init__(
+        self, 
+        root: str,      # dataset dir
+        meta: str,      # metadata
+        train: bool,    # training mode
+        length: int,    # video length
+        label: str,     # label field name
+        artery: str,    # left or right artery
+        inference: bool = False,
+        validation: bool = False,
+        transform: Optional[Callable] = None
+
+    ) -> None:
+        self.root = root
+        self.train = train
+        self.length = length
+        self.label = label
+        self.artery = artery
+        self.inference = inference
+        self.transform = transform
+        self.validation = validation
+        meta_path = meta if os.path.isabs(meta) else os.path.join(root, meta)
+
+        with open(meta_path) as f:
+            dataset = json.load(f)
+
+        if not self.inference:    
+            dataset = [rec for rec in dataset if len(rec[f"videos_{artery}"]) > 0]
+
+        if validation:
+            dataset = [rec for rec in dataset if rec[self.label] > 0]
+
+        self.dataset = dataset
+
+        artery_bin = {"left":0, "right":1}.get(artery.lower())
+        if artery_bin is None:
+            raise ValueError(f"Unknown artery '{artery}'")
+        
+        self.artery_bin = artery_bin
+
+    def __len__(self):
+        return len(self.dataset)
+    
+
+    def get_sample_weights(self):
+        # пороги для левой (0) и правой (1) артерии
+        bin_thresholds = {
+            0: [0, 5, 10, 15],   # левая
+            1: [0, 2, 5, 8],   # правая
+        }
+
+        # выберем пороги для текущей артерии
+        thresholds = bin_thresholds[self.artery_bin]
+
+        thr0, thr1, thr2, thr3 = thresholds
+
+        # разбиваем датасет по интервалам
+        self.dataset_0 = [rec for rec in self.dataset if rec[self.label] == thr0]
+        self.dataset_1 = [rec for rec in self.dataset if thr0 < rec[self.label] <= thr1]
+        self.dataset_2 = [rec for rec in self.dataset if thr1 < rec[self.label] <= thr2]
+        self.dataset_3 = [rec for rec in self.dataset if thr2 < rec[self.label] <= thr3]
+        self.dataset_4 = [rec for rec in self.dataset if rec[self.label] > thr3]
+
+
+        total = len(self.dataset_0) + len(self.dataset_1) + len(self.dataset_2) + len(self.dataset_3) + len(self.dataset_4)
+
+
+        def safe_weight(count):
+            return total / count if count > 0 else 0.0
+
+        self.weights_0 = safe_weight(len(self.dataset_0))
+        self.weights_1 = safe_weight(len(self.dataset_1))
+        self.weights_2 = safe_weight(len(self.dataset_2))
+        self.weights_3 = safe_weight(len(self.dataset_3))
+        self.weights_4 = safe_weight(len(self.dataset_4))
+
+        # print("Weights: ", self.weights_0, self.weights_1, self.weights_2, self.weights_3, self.weights_4)
+        print("Counts: ", len(self.dataset_0), len(self.dataset_1), len(self.dataset_2), len(self.dataset_3), len(self.dataset_4))
+
+        weights = []
+        for rec in self.dataset:
+            syntax_score = rec[self.label]
+            if syntax_score == thr0:
+                weights.append(self.weights_0)
+            elif thr0 < syntax_score <= thr1:
+                weights.append(self.weights_1)
+            elif thr1 < syntax_score <= thr2:
+                weights.append(self.weights_2)
+            elif thr2 < syntax_score <= thr3:
+                weights.append(self.weights_3)
+            else:
+                weights.append(self.weights_4)
+
+        self.weights = torch.tensor(weights, dtype=DTYPE)
+        return self.weights
+
+    def __getitem__(self, idx: int) -> Tuple[Tensor, int]:
+
+        rec = self.dataset[idx]
+        suid = rec["study_uid"]
+        
+         
+        if self.label:
+            bin_thresholds = {
+                0: 15,   # левая
+                1: 5,   # правая
+                }
+
+            label = torch.tensor([int(rec[self.label] > bin_thresholds[self.artery_bin])], dtype=DTYPE)
+            target = torch.tensor([np.log(1.0+rec[self.label])], dtype=DTYPE)
+        else:
+            label = torch.tensor([0], dtype=DTYPE)
+            target = torch.tensor([0], dtype=DTYPE)
+
+        nv = len(rec[f"videos_{self.artery}"])
+        if self.inference:
+            if nv == 0:
+                return 0, label, target, suid
+            seq = range(nv)
+        else:
+            seq = torch.randint(low=0, high=nv, size = (4,))
+
+        videos = []
+        for vi in seq:
+            video_rec = rec[f"videos_{self.artery}"][vi]
+            path = video_rec["path"]
+            if os.path.isabs(path):
+                full_path = path
+            else:
+                full_path = os.path.join(self.root, path)
+
+            video = pydicom.dcmread(full_path).pixel_array # Time, HW or WH
+
+            if video.dtype == np.uint16:
+                vmax = np.max(video)
+                assert vmax > 0
+                video = video.astype(np.float32)
+                video = video * (255. / vmax)
+                video = video.astype(np.uint8)
+            assert video.dtype == np.uint8
+
+            while len(video) < self.length:
+                video = np.concatenate([video, video])
+            t = len(video)
+            if self.train:
+                begin = torch.randint(low=0, high=t-self.length+1, size=(1,))
+                end = begin + self.length
+                video = video[begin:end, :, :]
+            else:
+                begin = (t - self.length) // 2
+                end = begin + self.length
+                video = video[begin:end, :, :]
+            
+            video = torch.tensor(np.stack([video, video, video], axis=-1))
+
+            if self.transform is not None:
+                video = self.transform(video)
+            videos.append(video)
+        videos = torch.stack(videos, dim=0)
+
+        
+        return videos, label, target, suid

full_model/rnn_model.py

@@ -0,0 +1,339 @@
+from typing import Any
+import torch
+import torch.nn.functional as F
+from torch import nn, optim
+import lightning.pytorch as pl
+import torchvision.models.video as tvmv
+import sklearn.metrics as skm
+import numpy as np
+
+
+class SyntaxLightningModule(pl.LightningModule):
+    """
+    Полная модель: 3D-ResNet backbone + RNN/Transformer head для SYNTAX score.
+
+    Варианты head (variant):
+    - mean_out: среднее по выходам backbone
+    - mean: среднее эмбеддингов + FC
+    - lstm_mean/lstm_last: LSTM (mean/last)
+    - gru_mean/gru_last: GRU (mean/last)
+    - bert_mean/bert_cls/bert_cls2: Transformer encoder
+    """
+
+    SUPPORTED_VARIANTS = [
+        "mean_out", "mean", "lstm_mean", "lstm_last",
+        "gru_mean", "gru_last", "bert_mean", "bert_cls", "bert_cls2"
+    ]
+
+    def __init__(
+        self,
+        num_classes: int,
+        lr: float,
+        variant: str,
+        weight_decay: float = 0.0,
+        max_epochs: int = None,
+        weight_path: str = None,        # путь к backbone-чекпоинту (.ckpt)
+        pl_weight_path: str = None,     # путь к полной модели (.ckpt или .pt)
+        pt_weights_format: bool = False, # True → .pt формат (torch.save), False → Lightning .ckpt
+        sigma_a: float = 0.0,
+        sigma_b: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        # Проверяем вариант head
+        if variant not in self.SUPPORTED_VARIANTS:
+            raise ValueError(f"variant must be one of {self.SUPPORTED_VARIANTS}")
+
+        self.num_classes = num_classes
+        self.variant = variant
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.max_epochs = max_epochs
+        self.sigma_a = sigma_a
+        self.sigma_b = sigma_b
+
+        # Backbone: 3D-ResNet
+        self.model = tvmv.r3d_18(weights=tvmv.R3D_18_Weights.DEFAULT)
+        in_features = self.model.fc.in_features
+
+        # Для большинства head заменяем fc на Identity (эмбеддинги)
+        if variant != "mean_out":
+            self.model.fc = nn.Identity()
+        else:
+            # mean_out использует финальные logits backbone
+            self.model.fc = nn.Linear(in_features, 2, bias=True)
+
+        # Загрузка backbone (если передан weight_path)
+        if weight_path is not None:
+            print(f"Loading backbone weights from {weight_path}")
+            self.load_weights_backbone(weight_path, self.model)
+
+        # Инициализация head в зависимости от variant
+        self._init_head(in_features, num_classes)
+
+        # Загрузка полной модели (если передан pl_weight_path)
+        if pl_weight_path is not None:
+            print(f"Loading full model weights from {pl_weight_path} (pt_format={pt_weights_format})")
+            self.load_full_model(pl_weight_path, pt_weights_format)
+
+        # Лоссы
+        self.loss_clf = nn.BCEWithLogitsLoss(reduction="none")
+        self.loss_reg = nn.MSELoss(reduction="none")
+
+        # Буферы метрик
+        self.y_val, self.p_val, self.r_val = [], [], []
+        self.ty_val, self.tp_val = [], []
+
+    def _init_head(self, in_features: int, num_classes: int):
+        """Инициализация head в зависимости от variant."""
+        if self.variant == "mean_out":
+            return  # используем self.model.fc
+
+        elif self.variant in ("gru_mean", "gru_last"):
+            self.rnn = nn.GRU(in_features, in_features // 4, batch_first=True)
+            self.dropout = nn.Dropout(0.2)
+            self.fc = nn.Linear(in_features // 4, num_classes, bias=True)
+
+        elif self.variant in ("lstm_mean", "lstm_last"):
+            self.lstm = nn.LSTM(
+                input_size=in_features,
+                hidden_size=in_features // 4,
+                proj_size=num_classes,
+                batch_first=True,
+            )
+
+        elif self.variant == "mean":
+            self.fc = nn.Linear(in_features, num_classes, bias=True)
+
+        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
+            encoder_layer = nn.TransformerEncoderLayer(
+                d_model=in_features,
+                nhead=4,
+                batch_first=True,
+                dim_feedforward=in_features // 4,
+            )
+            self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)
+            self.dropout = nn.Dropout(0.2)
+            self.fc = nn.Linear(in_features, num_classes, bias=True)
+            if self.variant == "bert_cls2":
+                self.cls = nn.Parameter(torch.randn(1, 1, in_features))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (batch, N_videos, C, T, H, W)
+        → (batch, N_videos, embed_dim) → head → (batch, num_classes)
+        """
+        batch_size, seq_len, *video_shape = x.shape
+        x = torch.flatten(x, start_dim=0, end_dim=1)  # (batch*seq, C, T, H, W)
+        x = self.model(x)  # (batch*seq, embed_dim)
+        x = torch.unflatten(x, 0, (batch_size, seq_len))  # (batch, seq, embed_dim)
+
+        # Head
+        if self.variant == "mean_out":
+            x = torch.mean(x, dim=1)  # mean по последовательности
+
+        elif self.variant in ("gru_mean", "gru_last"):
+            all_outs, last_out = self.rnn(x)
+            x = torch.mean(all_outs, dim=1) if self.variant == "gru_mean" else last_out
+            x = self.dropout(x)
+            x = self.fc(x)
+
+        elif self.variant in ("lstm_mean", "lstm_last"):
+            all_outs, (last_out, _) = self.lstm(x)
+            x = torch.mean(all_outs, dim=1) if self.variant == "lstm_mean" else last_out
+
+        elif self.variant == "mean":
+            x = torch.mean(x, dim=1)
+            x = self.fc(x)
+
+        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
+            if self.variant == "bert_cls":
+                x = F.pad(x, (0, 0, 1, 0), "constant", 0)  # prepend CLS
+            elif self.variant == "bert_cls2":
+                bs = x.size(0)
+                x = torch.cat([self.cls.expand(bs, -1, -1), x], dim=1)
+            x = self.encoder(x)
+            x = torch.mean(x, dim=1) if self.variant == "bert_mean" else x[:, 0, :]
+            x = self.dropout(x)
+            x = self.fc(x)
+
+        return x
+
+    def training_step(self, batch, batch_idx):
+        x, y, target, path = batch
+        y_hat = self(x)
+        yp_clf, yp_reg = y_hat[:, 0:1], y_hat[:, 1:]
+
+        # BCE с down-weight для отрицательных примеров
+        weights_clf = torch.where(y > 0, 1.0, 0.45)
+        clf_loss = (self.loss_clf(yp_clf, y) * weights_clf).mean()
+
+        # Регрессия с вариабельностью
+        reg_loss_raw = self.loss_reg(yp_reg, target)
+        sigma = self.sigma_a * target + self.sigma_b
+        reg_loss = (reg_loss_raw / (sigma ** 2)).mean()
+
+        loss = clf_loss + 0.5 * reg_loss
+
+        # Логирование
+        y_pred = torch.sigmoid(yp_clf)
+        y_bin = torch.round(y.detach().cpu()).int()
+        y_pred_bin = torch.round(y_pred.detach().cpu()).int()
+
+        self.log("train_clf_loss", clf_loss, prog_bar=True, sync_dist=True)
+        self.log("train_val_loss", reg_loss, prog_bar=True, sync_dist=True)
+        self.log("train_full_loss", loss, prog_bar=True, sync_dist=True)
+        self.log("train_f1", skm.f1_score(y_bin, y_pred_bin, zero_division=0),
+                 prog_bar=True, sync_dist=True)
+        self.log("train_acc", skm.accuracy_score(y_bin, y_pred_bin),
+                 prog_bar=True, sync_dist=True)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y, target, path = batch
+        y_hat = self(x)
+        yp_clf, yp_reg = y_hat[:, 0:1], y_hat[:, 1:]
+
+        # Аккумулируем для метрик
+        y_pred = torch.sigmoid(yp_clf)
+        self.y_val.append(int(y[..., 0].cpu()))
+        self.p_val.append(float(y_pred[..., 0].cpu()))
+        self.r_val.append(round(float(y_pred[..., 0].cpu())))
+        self.ty_val.append(float(target[..., 0].cpu()))
+        self.tp_val.append(float(yp_reg[..., 0].cpu()))
+
+        # Лосс (тот же, что и в train)
+        clf_loss = self.loss_clf(yp_clf, y).mean()
+        reg_loss_raw = self.loss_reg(yp_reg, target)
+        sigma = self.sigma_a * target + self.sigma_b
+        reg_loss = (reg_loss_raw / (sigma ** 2)).mean()
+        loss = clf_loss + 0.5 * reg_loss
+
+        return loss
+
+    def on_validation_epoch_end(self):
+        try:
+            auc = skm.roc_auc_score(self.y_val, self.p_val)
+            f1 = skm.f1_score(self.y_val, self.r_val, zero_division=0)
+            acc = skm.accuracy_score(self.y_val, self.r_val)
+            mae = skm.mean_absolute_error(self.y_val, self.r_val)
+            rmse = skm.root_mean_squared_error(self.ty_val, self.tp_val)
+
+            self.log("val_auc", auc, prog_bar=True, sync_dist=True)
+            self.log("val_f1", f1, prog_bar=True, sync_dist=True)
+            self.log("val_acc", acc, prog_bar=True, sync_dist=True)
+            self.log("val_mae", mae, prog_bar=True, sync_dist=True)
+            self.log("val_rmse", rmse, prog_bar=True, sync_dist=True)
+
+        except ValueError as err:
+            print(err)
+            print("Y_VAL", self.y_val[:10], "...")
+            print("P_VAL", self.p_val[:10], "...")
+
+        # Очистка буферов
+        [buf.clear() for buf in [self.y_val, self.p_val, self.r_val, self.ty_val, self.tp_val]]
+
+    def on_train_epoch_end(self):
+        lr = self.optimizers().optimizer.param_groups[0]["lr"]
+        self.log("lr", lr, on_epoch=True, sync_dist=True)
+
+    def configure_optimizers(self):
+        # Pretrain (заморозка backbone) или full fine-tune
+        if self.weight_path:
+            # Pretrain: обучаем только head
+            trainable_modules = self._get_trainable_head_modules()
+            for param in self.parameters():
+                param.requires_grad = False
+            for module in trainable_modules:
+                for param in module.parameters():
+                    param.requires_grad = True
+            params = [p for module in trainable_modules for p in module.parameters()]
+        else:
+            # Full: всё
+            for param in self.parameters():
+                param.requires_grad = True
+            params = self.parameters()
+
+        optimizer = optim.Adam(params, lr=self.lr, weight_decay=self.weight_decay)
+        if self.max_epochs:
+            scheduler = optim.lr_scheduler.OneCycleLR(
+                optimizer, max_lr=self.lr, total_steps=self.max_epochs
+            )
+            return [optimizer], [scheduler]
+        return optimizer
+
+    def _get_trainable_head_modules(self):
+        """Возвращает список обучаемых модулей head."""
+        if self.variant == "mean_out":
+            return [self.model.fc]
+        elif self.variant in ("gru_mean", "gru_last"):
+            return [self.rnn, self.fc]
+        elif self.variant in ("lstm_mean", "lstm_last"):
+            return [self.lstm]
+        elif self.variant == "mean":
+            return [self.fc]
+        elif self.variant in ("bert_mean", "bert_cls", "bert_cls2"):
+            modules = [self.encoder, self.fc]
+            if self.variant == "bert_cls2":
+                modules.append(self.cls)
+            return modules
+        return []
+
+    def load_weights_backbone(self, weight_path: str, model):
+        """Загрузка backbone из Lightning .ckpt."""
+        ckpt = torch.load(weight_path, map_location="cpu", weights_only=False)
+        state_dict = ckpt["state_dict"]
+        new_state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()}
+        model.load_state_dict(new_state_dict, strict=False)
+
+    def load_full_model(self, pl_weight_path: str, pt_weights_format: bool):
+        """Загрузка полной модели (.ckpt или .pt)."""
+        if pt_weights_format:
+            # .pt формат (torch.save)
+            state_dict = torch.load(pl_weight_path, map_location="cpu", weights_only=False)
+        else:
+            # Lightning .ckpt
+            ckpt = torch.load(pl_weight_path, map_location="cpu", weights_only=False)
+            state_dict = ckpt["state_dict"]
+
+        # Backbone
+        self.load_weights(state_dict, self.model, "model")
+
+        # Head
+        trainable_modules = self._get_trainable_head_modules()
+        for module in trainable_modules:
+            prefix = module.__class__.__name__.lower()
+            self.load_weights(state_dict, module, prefix)
+
+        if self.variant == "bert_cls2":
+            if "cls" in state_dict:
+                self.cls.data.copy_(state_dict["cls"])
+
+    def load_weights(self, state_dict, module, prefix: str):
+        """Загрузка весов модуля по префиксу."""
+        module_state = {
+            k.replace(f"{prefix}.", ""): v
+            for k, v in state_dict.items()
+            if k.startswith(prefix)
+        }
+        missing, unexpected = module.load_state_dict(module_state, strict=False)
+        if missing:
+            print(f"Missing keys for {prefix}: {len(missing)}")
+        if unexpected:
+            print(f"Unexpected keys for {prefix}: {len(unexpected)}")
+
+    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
+        x, y, target, path = batch
+        y_hat = self(x)
+        yp_clf, yp_reg = y_hat[:, 0:1], y_hat[:, 1:]
+        y_prob = torch.sigmoid(yp_clf)
+        return {
+            "y": y,
+            "y_pred": torch.round(y_prob),
+            "y_prob": y_prob,
+            "y_reg": yp_reg,
+            "target": target,
+        }

full_model/rnn_train.py

@@ -0,0 +1,206 @@
+import os
+import json
+import torch
+import numpy as np
+import click
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import TensorBoardLogger
+from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
+
+from pytorchvideo.transforms import Normalize, Permute, RandAugment
+from torch.utils.data import DataLoader
+from torchvision.transforms import transforms as T
+from torchvision.transforms._transforms_video import ToTensorVideo
+from torchvision.transforms import InterpolationMode
+
+from rnn_dataset import SyntaxDataset
+from rnn_model import SyntaxLightningModule
+
+torch.set_float32_matmul_precision("medium")
+
+
+"""
+Обучение RNN-head поверх предобученного backbone для SYNTAX score.
+
+Этапы:
+1) pretrain — обучается только head (backbone заморожен);
+2) full — fine-tuning всей модели (backbone + head).
+"""
+
+
+def get_transforms(video_size, imagenet_mean, imagenet_std, train=True):
+    """Трансформации для видео (train с аугментациями, test без)."""
+    interpolation_choices = [InterpolationMode.BILINEAR, InterpolationMode.BICUBIC]
+    if train:
+        return T.Compose([
+            ToTensorVideo(),
+            Permute(dims=[1, 0, 2, 3]),  # C,T,H,W → T,C,H,W
+            RandAugment(magnitude=10, num_layers=2),
+            T.RandomHorizontalFlip(),
+            Permute(dims=[1, 0, 2, 3]),  # T,C,H,W → C,T,H,W
+            T.RandomChoice([
+                T.Resize(size=video_size, interpolation=interp, antialias=True)
+                for interp in interpolation_choices
+            ]),
+            Normalize(mean=imagenet_mean, std=imagenet_std),
+        ])
+    return T.Compose([
+        ToTensorVideo(),
+        T.Resize(size=video_size, interpolation=InterpolationMode.BICUBIC, antialias=True),
+        Normalize(mean=imagenet_mean, std=imagenet_std),
+    ])
+
+
+def make_dataloader(dataset, batch_size, num_workers):
+    """DataLoader с shuffle (sampler закомментирован)."""
+    # dataset.get_sample_weights()  # можно включить WeightedRandomSampler
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        shuffle=True if not dataset.inference else False,
+        drop_last=True,
+        pin_memory=True,
+    )
+
+
+def make_model(num_classes, video_shape, lr, variant, weight_decay, max_epochs,
+               weight_path=None, pl_weight_path=None, pt_weights_format=False):
+    """Создание SyntaxLightningModule."""
+    return SyntaxLightningModule(
+        num_classes=num_classes,
+        lr=lr,
+        variant=variant,
+        weight_decay=weight_decay,
+        max_epochs=max_epochs,
+        weight_path=weight_path,
+        pl_weight_path=pl_weight_path,
+        pt_weights_format=pt_weights_format,
+    )
+
+
+def make_callbacks(artery: str, fold: int, phase: str):
+    """Callbacks: LR monitor + checkpoint по val_mae."""
+    lr_monitor = LearningRateMonitor(logging_interval="epoch")
+    if phase == "pre":
+        checkpoint = ModelCheckpoint(
+            monitor="val_mae",
+            save_top_k=1,
+            mode="min",
+            filename="model-{epoch:02d}-{val_rmse:.3f}",
+            save_last=True,
+        )
+    elif phase == "full":
+        checkpoint = ModelCheckpoint(
+            monitor="val_mae",
+            save_top_k=3,
+            mode="min",
+            filename="model-{epoch:02d}-{val_rmse:.3f}",
+            save_last=True,
+        )
+    else:
+        raise ValueError(f"phase must be 'pre' or 'full'")
+    return [lr_monitor, checkpoint]
+
+
+def make_trainer(max_epochs, logger_name, callbacks):
+    """Lightning Trainer с TensorBoard."""
+    logger = TensorBoardLogger(save_dir="rnn_logs", name=logger_name)
+    trainer = pl.Trainer(
+        max_epochs=max_epochs,
+        accelerator="gpu",
+        devices=1,  # легко поменять
+        strategy="ddp_find_unused_parameters_true",
+        precision="bf16-mixed",
+        callbacks=callbacks,
+        log_every_n_steps=10,
+        logger=logger,
+    )
+    return trainer
+
+
+@click.command()
+@click.option(
+    "-r", "--dataset-root", type=click.Path(exists=True), required=True,
+    help="Корень датасета (где лежат folds/*.json и DICOM).",
+)
+@click.option("--fold", type=int, default=0, help="Номер фолда (0-4).")
+@click.option("-a", "--artery", type=str, default="right", help="'left' или 'right'.")
+@click.option("--variant", type=str, default="lstm_mean", help="Тип head (lstm_mean и др.).")
+@click.option("-nc", "--num-classes", type=int, default=2)
+@click.option("-b", "--batch-size", type=int, default=8)
+@click.option("-f", "--frames-per-clip", type=int, default=32)
+@click.option("-v", "--video-size", type=click.Tuple([int, int]), default=(256, 256))
+@click.option("--max-epochs", type=int, default=10)
+@click.option("--num-workers", type=int, default=8)
+@click.option("--fast-dev-run", is_flag=True)
+@click.option("--seed", type=int, default=42)
+@click.option("--backbone-ckpt", type=str, help="Путь к backbone-чекпоинту для pretrain.")
+def main(
+    dataset_root, fold, artery, variant, num_classes, batch_size, frames_per_clip,
+    video_size, max_epochs, num_workers, fast_dev_run, seed, backbone_ckpt,
+):
+    pl.seed_everything(seed)
+    artery = artery.lower()
+    artery_bin = {"left": 0, "right": 1}[artery]
+
+    print(f"Training {variant} head for {artery} artery, fold {fold}")
+
+    imagenet_mean = [0.485, 0.456, 0.406]
+    imagenet_std = [0.229, 0.224, 0.225]
+
+    # Datasets с относительными путями
+    train_meta = os.path.join("rnn_folds", f"step2_rnn_fold{fold:02d}_train.json")
+    val_meta = os.path.join("rnn_folds", f"step2_rnn_fold{fold:02d}_eval.json")
+
+    train_set = SyntaxDataset(
+        root=dataset_root,
+        meta=train_meta,
+        train=True,
+        length=frames_per_clip,
+        label=f"syntax_{artery}",
+        artery=artery,
+        transform=get_transforms(video_size, imagenet_mean, imagenet_std, train=True),
+    )
+    val_set = SyntaxDataset(
+        root=dataset_root,
+        meta=val_meta,
+        train=False,
+        length=frames_per_clip,
+        label=f"syntax_{artery}",
+        artery=artery,
+        validation=True,
+        transform=get_transforms(video_size, imagenet_mean, imagenet_std, train=False),
+    )
+
+    # DataLoaders
+    train_loader_pre = make_dataloader(train_set, batch_size * 2, num_workers)
+    train_loader_post = make_dataloader(train_set, batch_size, num_workers)
+    val_loader = make_dataloader(val_set, 1, num_workers)
+
+    # Форма видео
+    x, *_ = next(iter(train_loader_pre))
+    video_shape = x.shape[1:]
+
+    # Pretrain head
+    callbacks_pre = make_callbacks(artery, fold, "pre")
+    model_pre = make_model(
+        num_classes, video_shape, lr=1e-4, variant=variant,
+        weight_decay=0.01, max_epochs=max_epochs, weight_path=backbone_ckpt,
+    )
+    trainer_pre = make_trainer(max_epochs, f"{artery}_{variant}_pre_fold{fold:02d}", callbacks_pre)
+    trainer_pre.fit(model_pre, train_loader_pre, val_loader)
+
+    # Full fine-tune
+    callbacks_full = make_callbacks(artery, fold, "full")
+    model_full = make_model(
+        num_classes, video_shape, lr=2e-5, variant=variant,
+        weight_decay=0.01, max_epochs=max_epochs,
+        pl_weight_path=trainer_pre.checkpoint_callback.best_model_path,
+    )
+    trainer_full = make_trainer(max_epochs, f"{artery}_{variant}_full_fold{fold:02d}", callbacks_full)
+    trainer_full.fit(model_full, train_loader_post, val_loader)
+
+
+if __name__ == "__main__":
+    main()