Spaces:
Running on Zero
Running on Zero
File size: 1,772 Bytes
093748c | 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 | from machine_learning.Dataset import BindingResidueDataset
from machine_learning.ModelManager import load_classifier_torch
from setup.configProcessor import get_weight_dir, get_structure_cutoff
from setup.generalSetup import select_device
from torch_geometric.loader import DataLoader
import os
import torch
import numpy as np
def run_prediction(model_parameters):
structure_cutoff = get_structure_cutoff(only_first_value=True)
dataset = BindingResidueDataset(structure_cutoff)
predictions = {}
sigmoid = torch.nn.Sigmoid()
models = os.listdir(get_weight_dir())
for model in models:
model_path = os.path.join(get_weight_dir(), model)
pretrained_model = load_classifier_torch(model_path, model_parameters)
pretrained_model.to(select_device())
data_loader = DataLoader(dataset, batch_size=1, shuffle=True)
pretrained_model.eval()
for protein_graph, protein_id in data_loader:
with torch.no_grad():
protein_id = protein_id[0]
protein_graph = protein_graph.to(select_device())
prediction = pretrained_model.forward(
protein_graph.x,
protein_graph.edge_index,
protein_graph.edge_index_cutoff,
protein_graph.edge_features,
protein_graph.dssp_features
)
prediction = sigmoid(prediction)
prediction = prediction.detach().cpu()
prediction = np.array(prediction)
if protein_id not in predictions:
predictions[protein_id] = [prediction]
else:
predictions[protein_id] += [prediction]
return predictions
|