|
import os |
|
import torch |
|
from torch.utils.data import Dataset |
|
from pathlib import Path |
|
import argparse |
|
from scripts.genomic_plip_model import GenomicPLIPModel |
|
from transformers import CLIPVisionModel |
|
|
|
class PatientTileDataset(Dataset): |
|
def __init__(self, data_dir, model, save_dir): |
|
super().__init__() |
|
self.data_dir = data_dir |
|
self.model = model |
|
self.save_dir = Path(save_dir) |
|
self.files = [] |
|
for patient_id in os.listdir(data_dir): |
|
patient_dir = os.path.join(data_dir, patient_id) |
|
if os.path.isdir(patient_dir): |
|
for f in os.listdir(patient_dir): |
|
if f.endswith('.pt'): |
|
self.files.append((os.path.join(patient_dir, f), patient_id)) |
|
|
|
def __len__(self): |
|
return len(self.files) |
|
|
|
def __getitem__(self, idx): |
|
file_path, patient_id = self.files[idx] |
|
data = torch.load(file_path) |
|
tile_data = torch.from_numpy(data['tile_data'][0]).unsqueeze(0) |
|
with torch.no_grad(): |
|
vision_features, _ = self.model(pixel_values=tile_data, score_vector=torch.zeros(1, 4)) |
|
feature_path = self.save_dir / patient_id / os.path.basename(file_path) |
|
feature_path.parent.mkdir(parents=True, exist_ok=True) |
|
torch.save(vision_features, feature_path) |
|
return feature_path |
|
|
|
def extract_features(data_dir, save_dir, model_path): |
|
original_model = CLIPVisionModel.from_pretrained("./plip/") |
|
custom_model = GenomicPLIPModel(original_model) |
|
custom_model.load_state_dict(torch.load(model_path)) |
|
custom_model.eval() |
|
|
|
dataset = PatientTileDataset(data_dir=data_dir, model=custom_model, save_dir=save_dir) |
|
for _ in dataset: |
|
pass |
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser(description="Extract features from genomic aligned tiles.") |
|
parser.add_argument('--data_dir', type=str, default='plip_preprocess/', help='Directory containing the pre processed patient data.') |
|
parser.add_argument('--save_dir', type=str, default='omics_align_features/', help='Directory to save the extracted features.') |
|
parser.add_argument('--model_path', type=str, default='./save_model/omics_plip.pth', help='Path to the trained model file.') |
|
|
|
args = parser.parse_args() |
|
|
|
extract_features(args.data_dir, args.save_dir, args.model_path) |
|
|
