add backbone model

backbone/dataset.py

@@ -0,0 +1,222 @@
+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
+
+# Полуточность достаточно для хранения весов и таргетов,
+# а сами вычисления в модели идут в float32 / bf16.
+DTYPE = torch.float16
+
+
+class SyntaxDataset(Dataset):
+    """
+    PyTorch Dataset для обучения видеобэкбона на задаче SYNTAX.
+
+    Функциональность:
+    - читает метаданные из JSON (относительный путь относительно root);
+    - фильтрует по артерии (левая / правая);
+    - опционально отфильтровывает только примеры с положительным SYNTAX
+      (validation=True);
+    - рассчитывает sample weights по бинам SYNTAX (для WeightedRandomSampler);
+    - конвертирует DICOM-видео в тензор (T, H, W, 3) c uint8 [0–255];
+    - возвращает:
+        video, label_bin, target_log, weight, rel_path, original_label.
+    """
+
+    def __init__(
+        self,
+        root: str,                      # корневая директория датасета
+        meta: str,                      # относительный путь к JSON с метаданными
+        train: bool,                    # режим: train / eval
+        length: int,                    # длина клипа (кол-во кадров)
+        label: str,                     # имя поля с SYNTAX score в JSON
+        artery_bin: int,                # 0 — левая, 1 — правая артерия
+        validation: bool = False,       # отбрасывать ли нулевые SYNTAX
+        transform: Optional[Callable] = None,
+    ) -> None:
+        super().__init__()
+        self.root = root
+        self.train = train
+        self.length = length
+        self.label = label
+        self.transform = transform
+        self.validation = validation
+
+        # meta теперь трактуется как ОТНОСИТЕЛЬНЫЙ путь от root
+        meta_path = os.path.join(root, meta)
+        with open(meta_path, "r") as f:
+            dataset = json.load(f)
+
+        # Фильтр по артерии (0 — левая, 1 — правая)
+        if artery_bin is not None:
+            assert artery_bin in (0, 1), "artery_bin должен быть 0 (левая) или 1 (правая)"
+            dataset = [rec for rec in dataset if rec["artery"] == artery_bin]
+            self.artery_bin = artery_bin
+        else:
+            # Для корректной работы get_sample_weights ожидаем известный artery_bin
+            raise ValueError("artery_bin должен быть явно задан (0 или 1).")
+
+        # Валидационный набор: берём только записи с положительным SYNTAX
+        if validation:
+            dataset = [rec for rec in dataset if rec[self.label] > 0]
+
+        # Инициализируем веса с единиц
+        for rec in dataset:
+            rec["weight"] = 1.0
+
+        self.dataset = dataset
+
+    # ------------------------------------------------------------------
+    # Веса для WeightedRandomSampler
+    # ------------------------------------------------------------------
+    def get_sample_weights(self) -> Tensor:
+        """
+        Считает веса для примеров по бинам SYNTAX.
+
+        Для каждой артерии определён свой набор порогов,
+        после чего каждый пример получает вес, обратный частоте своего бина.
+        """
+        # Пороговые значения по артериям (подбирались эмпирически)
+        bin_thresholds = {
+            0: [0, 5, 10, 15],  # левая
+            1: [0, 2, 5, 8],    # правая
+        }
+
+        thresholds = bin_thresholds[self.artery_bin]
+        thr0, thr1, thr2, thr3 = thresholds
+
+        # Бины по значениям SYNTAX
+        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: int) -> float:
+            # Если в би��е нет примеров, вес ставим 0.0
+            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 __len__(self) -> int:
+        return len(self.dataset)
+
+    # ------------------------------------------------------------------
+    def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, str, Tensor]:
+        """
+        Возвращает один пример:
+        - video: Tensor (T, H, W, 3) → после transform обычно (C, T, H, W)
+        - label: бинарный таргет по порогу для конкретной артерии
+        - target: логарифмированный SYNTAX score (регрессия)
+        - weight: вес примера (для самплера / лосса)
+        - path: относительный путь к DICOM файлу
+        - original_label: исходный SYNTAX score
+        """
+        rec = self.dataset[idx]
+
+        # Относительный путь к DICOM из JSON (мы не храним абсолютные пути)
+        path = rec["path"]
+        weight = rec["weight"]
+
+        full_path = os.path.join(self.root, path)
+        video = pydicom.dcmread(full_path).pixel_array  # (T, H, W)
+
+        # Приводим 16-битный сигнал к диапазону [0, 255] uint8
+        if video.dtype == np.uint16:
+            vmax = np.max(video)
+            assert vmax > 0
+            video = video.astype(np.float32)
+            video = video * (255.0 / vmax)
+            video = video.astype(np.uint8)
+        assert video.dtype == np.uint8
+
+        # Порог для бинарной классификации зависит от артерии
+        bin_thresholds = {
+            0: 15,  # левая
+            1: 5,   # правая
+        }
+
+        syntax_value = rec[self.label]
+        label = torch.tensor(
+            [int(syntax_value > bin_thresholds[self.artery_bin])],
+            dtype=DTYPE,
+        )
+        target = torch.tensor([np.log(1.0 + syntax_value)], dtype=DTYPE)
+        original_label = torch.tensor([syntax_value], dtype=DTYPE)
+
+        # Дублируем видео по времени, пока не наберём нужную длину клипа
+        while len(video) < self.length:
+            video = np.concatenate([video, video])
+        t = len(video)
+
+        if self.train:
+            # Случайный подотрезок длины self.length
+            begin = torch.randint(low=0, high=t - self.length + 1, size=(1,))
+            end = begin + self.length
+            video = video[begin:end, :, :]
+        else:
+            # В валидации используем весь видеоряд (обрежется трансформами / моделью)
+            video = video
+
+        # Превращаем (T, H, W) → (T, H, W, 3) путём копирования каналов (grayscale→RGB)
+        video = torch.tensor(np.stack([video, video, video], axis=-1))
+
+        if self.transform is not None:
+            video = self.transform(video)
+
+        sample_weight = torch.tensor([weight], dtype=DTYPE)
+
+        return video, label, target, sample_weight, path, original_label

backbone/pl_model.py

@@ -0,0 +1,244 @@
+from typing import Any
+import torch
+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):
+    """
+    LightningModule для обучения 3D-ResNet (r3d_18) как backbone
+    в задаче предсказания SYNTAX score по видеоангиографии.
+
+    Модель предсказывает:
+      - yp_clf: вероятность поражения (syntax > порог) — бинарная классификация
+      - yp_reg: логарифмированное значение SYNTAX — регрессия
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        lr: float,
+        weight_decay: float = 0.0,
+        max_epochs: int = None,
+        weight_path: str = None,
+        sigma_a: float = 0.0,
+        sigma_b: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.num_classes = num_classes
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.max_epochs = max_epochs
+        self.weight_path = weight_path
+        self.sigma_a = sigma_a
+        self.sigma_b = sigma_b
+
+        # Базовый 3D-ResNet с ImageNet Kinetics-предобученными весами
+        self.model = tvmv.r3d_18(weights=tvmv.R3D_18_Weights.DEFAULT)
+
+        # Последний слой заменяем на Linear с num_classes выходами:
+        #   1 канал для классификации, 1 для регрессии
+        in_features = self.model.fc.in_features
+        self.model.fc = nn.Linear(in_features=in_features, out_features=num_classes, bias=True)
+
+        # Если передан путь к чекпоинту Lightning — загружаем backbone
+        if self.weight_path is not None:
+            ckpt = torch.load(self.weight_path, map_location="cpu", weights_only=False)
+            state_dict = ckpt["state_dict"]
+            # Чистим префикс "model." у ключей
+            new_state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()}
+            self.model.load_state_dict(new_state_dict, strict=False)
+
+        # Лоссы
+        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 forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+    # ------------------------------------------------------------------
+    def training_step(self, batch, batch_idx):
+        """
+        Один шаг обучения:
+        - бинарная классификация поражения (BCE с down-weight для нулей);
+        - регрессия логарифмированного SYNTAX с учётом get_sigma(target).
+        """
+        x, y, target, sample_weight, path, original_label = batch
+
+        y_hat = self(x)
+        yp_clf = y_hat[:, 0:1]   # logits для классификации
+        yp_reg = y_hat[:, 1:]    # регрессия (лог SYNTAX)
+
+        # BCE с меньшим весом для класса 0 (нет поражения)
+        weights_clf = torch.where(y > 0, 1.0, 0.45)
+        clf_loss = self.loss_clf(yp_clf, y)
+        clf_loss = (clf_loss * 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, sample_weight, path, original_label = batch
+
+        y_hat = self(x)
+        yp_clf = y_hat[:, 0:1]
+        yp_reg = y_hat[:, 1:]
+
+        # Комбинированный лосс
+        clf_loss = self.loss_clf(yp_clf, y)
+        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.mean() + 0.5 * reg_loss
+
+        # Для метрик
+        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()))
+
+        return loss
+
+    # ------------------------------------------------------------------
+    def on_validation_epoch_end(self) -> None:
+        """
+        Подсчёт валидационных метрик по всей эпохе и логирование в Logger.
+        """
+        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)
+            print("P_VAL", self.p_val)
+
+        # Чистим буферы к следующей эпохе
+        self.y_val.clear()
+        self.p_val.clear()
+        self.r_val.clear()
+        self.ty_val.clear()
+        self.tp_val.clear()
+
+    # ------------------------------------------------------------------
+    def on_train_epoch_end(self) -> None:
+        """Логирование текущего learning rate."""
+        opt = self.optimizers()
+        if hasattr(opt, "optimizer"):
+            lr = opt.optimizer.param_groups[0]["lr"]
+        else:
+            lr = opt.param_groups[0]["lr"]
+        self.log("lr", lr, on_step=False, on_epoch=True, sync_dist=True)
+
+    # ------------------------------------------------------------------
+    def configure_optimizers(self):
+        """
+        - Если weight_path не задан → pretrain: обучаем только финальный fc-слой.
+        - Если weight_path задан → full fine-tuning: обучаем весь backbone.
+        """
+        if not self.weight_path:
+            # Pretrain: замораживаем всё, кроме финального слоя
+            for param in self.parameters():
+                param.requires_grad = False
+            for p in self.model.fc.parameters():
+                p.requires_grad = True
+            params = list(self.model.fc.parameters())
+        else:
+            # Full fine-tune: обучаем все параметры модели
+            for param in self.parameters():
+                param.requires_grad = True
+            params = self.parameters()
+
+        optimizer = optim.AdamW(params, lr=self.lr, weight_decay=self.weight_decay)
+
+        if self.max_epochs is not None:
+            scheduler = optim.lr_scheduler.OneCycleLR(
+                optimizer=optimizer,
+                max_lr=self.lr,
+                total_steps=self.max_epochs,
+            )
+            return [optimizer], [scheduler]
+        else:
+            return optimizer
+
+    # ------------------------------------------------------------------
+    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
+        """
+        Инференс: возвращает словарь с бинарным предсказанием, вероятностями
+        и регрессионным выходом.
+        """
+        x, y, target, sample_weight, path, original_label = batch
+        y_hat = self(x)
+        yp_clf = y_hat[:, 0:1]
+        yp_reg = y_hat[:, 1:]
+        y_prob = torch.sigmoid(yp_clf)
+        y_pred = torch.round(y_prob)
+
+        return {
+            "y": y,
+            "y_pred": y_pred,
+            "y_prob": y_prob,
+            "y_reg": yp_reg,
+            "target": target,
+            "original_label": original_label,
+        }

backbone/pl_train.py

@@ -0,0 +1,278 @@
+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 lightning.pytorch.profilers import AdvancedProfiler, PyTorchProfiler
+
+from pytorchvideo.transforms import Normalize, Permute, RandAugment
+from torch.utils.data import DataLoader, WeightedRandomSampler
+from torchvision.transforms import transforms as T
+from torchvision.transforms._transforms_video import ToTensorVideo
+from torchvision.transforms import InterpolationMode
+
+from dataset import SyntaxDataset
+from pl_model import SyntaxLightningModule
+
+import warnings
+warnings.filterwarnings("ignore", message="No device id is provided via `init_process_group`")
+
+torch.set_float32_matmul_precision("medium")
+
+
+"""
+Скрипт обучения backbone (3D-ResNet) для предсказания SYNTAX score.
+
+Шаги:
+1) предварительное обучение (pretrain) — обучается только последний слой;
+2) полное дообучение (full) — fine-tuning всего backbone.
+"""
+
+
+# ------------------- Трансформации -------------------
+def get_transforms(video_size, imagenet_mean, imagenet_std, train=True):
+    interpolation_choices = [
+        InterpolationMode.BILINEAR,
+        InterpolationMode.BICUBIC,
+    ]
+    if train:
+        return T.Compose([
+            ToTensorVideo(),                        # (T, H, W, 3) -> (C, T, H, W)
+            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),
+        ])
+    else:
+        return T.Compose([
+            ToTensorVideo(),
+            T.Resize(size=video_size, interpolation=InterpolationMode.BICUBIC, antialias=True),
+            Normalize(mean=imagenet_mean, std=imagenet_std),
+        ])
+
+
+# ------------------- DataLoader -------------------
+def make_dataloader(dataset, batch_size, num_workers):
+    """
+    Создаёт DataLoader; по умолчанию используем shuffle,
+    но можно легко переключиться на WeightedRandomSampler.
+    """
+    sample_weights = dataset.get_sample_weights()
+    # sampler = WeightedRandomSampler(sample_weights, len(dataset), replacement=True)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        # sampler=sampler,
+        shuffle=True,
+        drop_last=True,
+        pin_memory=True,
+    )
+
+
+# ------------------- Модель -------------------
+def make_model(num_classes, video_shape, lr, weight_decay, max_epochs, weight_path=None):
+    """
+    Обёртка над SyntaxLightningModule для единообразного создания модели
+    на этапах pretrain и full fine-tuning.
+    """
+    model = SyntaxLightningModule(
+        num_classes=num_classes,
+        lr=lr,
+        weight_decay=weight_decay,
+        max_epochs=max_epochs,
+        weight_path=weight_path,
+    )
+    return model
+
+
+# ------------------- Callbacks -------------------
+def make_callbacks(artery: str, fold: int, phase: str):
+    """
+    Возвращает набор callback'ов:
+    - LearningRateMonitor
+    - ModelCheckpoint с сохранением по наилучшему 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"Unknown phase '{phase}', expected 'pre' or 'full'")
+
+    return [lr_monitor, checkpoint]
+
+
+# ------------------- Trainer -------------------
+def make_trainer(max_epochs, logger_name, callbacks):
+    """
+    Создаёт Lightning Trainer c TensorBoardLogger.
+
+    Важно: пути к логам и устройствам можно адаптировать под свой кластер.
+    """
+    logger = TensorBoardLogger(
+        save_dir="backbone_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),
+    default=".",
+    required=True,
+    help="Путь к корню датасета (директория, внутри которой лежат JSON и DICOM).",
+)
+@click.option("--fold", type=int, default=0, required=True, help="Номер фолда (0–4).")
+@click.option(
+    "-a",
+    "--artery",
+    type=str,
+    default="right",
+    required=True,
+    help="Название артерии: 'left' или 'right'.",
+)
+@click.option("-nc", "--num-classes", type=int, default=2, help="Число выходных каналов модели.")
+@click.option("-b", "--batch-size", type=int, default=50, help="Размер batch.")
+@click.option("-f", "--frames-per-clip", type=int, default=32, help="Количество кадров в клипе.")
+@click.option(
+    "-v",
+    "--video-size",
+    type=click.Tuple([int, int]),
+    default=(256, 256),
+    help="Размер кадра (H, W).",
+)
+@click.option("--max-epochs", type=int, default=10, help="Число эпох на этапе full fine-tuning.")
+@click.option("--num-workers", type=int, default=8, help="Число воркеров для DataLoader.")
+@click.option(
+    "--fast-dev-run",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    help="Режим быстрой проверки пайплайна (1–2 батча).",
+)
+@click.option("--seed", type=int, default=42, help="Сид для воспроизводимости.")
+def main(
+    dataset_root,
+    fold,
+    artery,
+    num_classes,
+    batch_size,
+    frames_per_clip,
+    video_size,
+    max_epochs,
+    num_workers,
+    fast_dev_run,
+    seed,
+):
+    pl.seed_everything(seed)
+
+    artery = artery.lower()
+    artery_bin = {"left": 0, "right": 1}.get(artery)
+    if artery_bin is None:
+        raise ValueError(f"Unknown artery '{artery}', expected 'left' or 'right'.")
+
+    imagenet_mean = [0.485, 0.456, 0.406]
+    imagenet_std = [0.229, 0.224, 0.225]
+
+    # ------------------- Datasets -------------------
+    # Путь к JSON теперь относительный относительно dataset_root
+    train_meta = os.path.join("folds", f"step2_fold{fold:02d}_train.json")
+    val_meta = os.path.join("folds", f"step2_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_bin=artery_bin,
+        validation=False,
+        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_bin=artery_bin,
+        validation=True,
+        transform=get_transforms(video_size, imagenet_mean, imagenet_std, train=False),
+    )
+
+    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)
+
+    # Получаем форму входного видео (C, T, H, W) из одного батча
+    x, *_ = next(iter(train_loader_pre))
+    video_shape = x.shape[1:]
+
+    # ------------------- Callbacks -------------------
+    callbacks_pre = make_callbacks(artery=artery, fold=fold, phase="pre")
+    callbacks_full = make_callbacks(artery=artery, fold=fold, phase="full")
+
+    # ------------------- Pretrain -------------------
+    num_pre_epochs = 10
+    model_pre = make_model(
+        num_classes=num_classes,
+        video_shape=video_shape,
+        lr=3e-4,
+        weight_decay=0.01,
+        max_epochs=num_pre_epochs,
+    )
+    trainer_pre = make_trainer(num_pre_epochs, f"{artery}BinSyntax_R3D_pre_fold{fold:02d}", callbacks_pre)
+    trainer_pre.fit(model_pre, train_loader_pre, val_loader, ckpt_path=None)
+
+    # ------------------- Full train -------------------
+    model_full = make_model(
+        num_classes=num_classes,
+        video_shape=video_shape,
+        lr=1e-4,
+        weight_decay=0.01,
+        max_epochs=max_epochs,
+        weight_path=trainer_pre.checkpoint_callback.last_model_path,
+    )
+    trainer_full = make_trainer(max_epochs, f"{artery}BinSyntax_R3D_full_fold{fold:02d}", callbacks_full)
+    trainer_full.fit(model_full, train_loader_post, val_loader, ckpt_path=None)
+
+
+if __name__ == "__main__":
+    main()