Spaces:
Runtime error
Runtime error
File size: 5,027 Bytes
54660f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
from cl import IClassifier
from build_graph_utils import *
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision.models as models
import torchvision.transforms.functional as VF
from torchvision import transforms
import sys, argparse, os, glob
import pandas as pd
import numpy as np
from PIL import Image
from collections import OrderedDict
def compute_feats(args, bags_list, i_classifier, device, save_path=None, whole_slide_path=None):
num_bags = len(bags_list)
Tensor = torch.FloatTensor
for i in range(0, num_bags):
feats_list = []
if args.magnification == '20x':
glob_path = os.path.join(bags_list[i], '*.jpeg')
csv_file_path = glob.glob(glob_path)
# line below was in the original version, commented due to errror with current version
#file_name = bags_list[i].split('/')[-3].split('_')[0]
file_name = glob_path.split('/')[-3].split('_')[0]
if args.magnification == '5x' or args.magnification == '10x':
csv_file_path = glob.glob(os.path.join(bags_list[i], '*.jpg'))
dataloader, bag_size = bag_dataset(args, csv_file_path)
print('{} files to be processed: {}'.format(len(csv_file_path), file_name))
if os.path.isdir(os.path.join(save_path, 'simclr_files', file_name)) or len(csv_file_path) < 1:
print('alreday exists')
continue
with torch.no_grad():
for iteration, batch in enumerate(dataloader):
patches = batch['input'].float().to(device)
feats, classes = i_classifier(patches)
#feats = feats.cpu().numpy()
feats_list.extend(feats)
os.makedirs(os.path.join(save_path, 'simclr_files', file_name), exist_ok=True)
txt_file = open(os.path.join(save_path, 'simclr_files', file_name, 'c_idx.txt'), "w+")
save_coords(txt_file, csv_file_path)
# save node features
output = torch.stack(feats_list, dim=0).to(device)
torch.save(output, os.path.join(save_path, 'simclr_files', file_name, 'features.pt'))
# save adjacent matrix
adj_s = adj_matrix(csv_file_path, output, device=device)
torch.save(adj_s, os.path.join(save_path, 'simclr_files', file_name, 'adj_s.pt'))
print('\r Computed: {}/{}'.format(i+1, num_bags))
def main():
parser = argparse.ArgumentParser(description='Compute TCGA features from SimCLR embedder')
parser.add_argument('--num_classes', default=2, type=int, help='Number of output classes')
parser.add_argument('--num_feats', default=512, type=int, help='Feature size')
parser.add_argument('--batch_size', default=128, type=int, help='Batch size of dataloader')
parser.add_argument('--num_workers', default=0, type=int, help='Number of threads for datalodaer')
parser.add_argument('--dataset', default=None, type=str, help='path to patches')
parser.add_argument('--backbone', default='resnet18', type=str, help='Embedder backbone')
parser.add_argument('--magnification', default='20x', type=str, help='Magnification to compute features')
parser.add_argument('--weights', default=None, type=str, help='path to the pretrained weights')
parser.add_argument('--output', default=None, type=str, help='path to the output graph folder')
args = parser.parse_args()
if args.backbone == 'resnet18':
resnet = models.resnet18(pretrained=False, norm_layer=nn.InstanceNorm2d)
num_feats = 512
if args.backbone == 'resnet34':
resnet = models.resnet34(pretrained=False, norm_layer=nn.InstanceNorm2d)
num_feats = 512
if args.backbone == 'resnet50':
resnet = models.resnet50(pretrained=False, norm_layer=nn.InstanceNorm2d)
num_feats = 2048
if args.backbone == 'resnet101':
resnet = models.resnet101(pretrained=False, norm_layer=nn.InstanceNorm2d)
num_feats = 2048
for param in resnet.parameters():
param.requires_grad = False
resnet.fc = nn.Identity()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Running on:", device)
i_classifier = IClassifier(resnet, num_feats, output_class=args.num_classes).to(device)
# load feature extractor
if args.weights is None:
print('No feature extractor')
return
state_dict_weights = torch.load(args.weights)
state_dict_init = i_classifier.state_dict()
new_state_dict = OrderedDict()
for (k, v), (k_0, v_0) in zip(state_dict_weights.items(), state_dict_init.items()):
if 'features' not in k:
continue
name = k_0
new_state_dict[name] = v
i_classifier.load_state_dict(new_state_dict, strict=False)
os.makedirs(args.output, exist_ok=True)
bags_list = glob.glob(args.dataset)
print(bags_list)
compute_feats(args, bags_list, i_classifier, device, args.output)
if __name__ == '__main__':
main()
|