feat: train fixed-size classifiers and export class probabilities

classify.py

@@ -16,6 +16,7 @@ import json
 from pathlib import Path
 
 from datasets import FolderClassificationDataset, classification_collate
+from datasets.tile_classification import YoloTileClassificationDataset
 from engine import ClassificationTrainer
 from models import build_model
 from utils import load_config, seed_everything, setup_logging
@@ -40,6 +41,13 @@ def build_loaders(cfg, split_cfg, classes=None):
 
     datasets = {}
     for split, transforms_key in (("train", "train_transforms"), ("val", "val_transforms")):
+        if split_cfg.get("format", "folder") == "yolo_tiles":
+            datasets[split] = YoloTileClassificationDataset(
+                yaml_path=split_cfg["yaml"], split=split,
+                image_size=split_cfg["image_size"], transforms=split_cfg[transforms_key],
+                classes=classes or split_cfg.get("classes"),
+            )
+            continue
         datasets[split] = FolderClassificationDataset(
             root=split_cfg["root"],
             split=split,