read_data.py

import os
import h5py
import torch
import pandas as pd
from torch.utils.data import Dataset
import numpy as np
import pdb

class SuperTileRNADataset(Dataset):
    def __init__(self, df, feature_path):
        self.df = df
        self.feature_path = feature_path

        # Infer number of genes (columns that start with 'rna_')
        self.rna_cols = [col for col in df.columns if col.startswith('rna_')]
        self.num_genes = len(self.rna_cols)

        # Infer feature dimension from the first sample
        first_file = df.iloc[0]['wsi_file_name']
        self.feature_dim = self._infer_feature_dim(first_file)

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        try:
            row = self.df.iloc[idx]
            wsi_file = row['wsi_file_name']
            wsi_base = os.path.splitext(os.path.basename(wsi_file))[0]

            feature_file = os.path.join(self.feature_path, wsi_base + '.h5')
            if not os.path.exists(feature_file):
                print(f"Feature file not found: {feature_file}")
                return None

            with h5py.File(feature_file, 'r') as f:
                if 'cluster_features' not in f:
                    print(f"No cluster_features found in: {feature_file}")
                    return None
                feats = f['cluster_features'][:]  # Shape (K=500, 2048)

            feats = torch.tensor(feats, dtype=torch.float32)
            rna = torch.tensor(row[self.rna_cols].values.astype(np.float32))

            return feats, rna, row['wsi_file_name'], row['tcga_project']

        except Exception as e:
            print(f"Error loading data for index {idx}: {e}")
            return None

    def _infer_feature_dim(self, wsi_file):
        try:
            wsi_base = os.path.splitext(os.path.basename(wsi_file))[0]
            path = os.path.join(self.feature_path, wsi_base + '.h5')
            with h5py.File(path, 'r') as f:
                return f['cluster_features'].shape[1]
        except Exception as e:
            print(f"Feature dim inference failed: {e}")
            return 2048  # fallback