classify.py

#!/usr/bin/env python3

import json
import sys

import joblib
import numpy as np
import torch

if len(sys.argv) < 3:
    print('Usage: %s <input BGC features file path> <output classification file path>' % sys.argv[0])
    exit(1)
input_path = sys.argv[1]
output_path = sys.argv[2]

with open(input_path, 'r') as f:
    feature_data = json.load(f)
X_names = sorted(feature_data.keys())
X = np.array([feature_data[bgc_name] for bgc_name in X_names])

scaler = joblib.load('data/class-scaler.pkl')
X = scaler.transform(X)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

X = torch.tensor(X, dtype=torch.float32).to(device)

class MLP(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super(MLP, self).__init__()
        layers = []
        sizes = [input_size // 2, input_size // 4, input_size // 8]
        for i, size in enumerate(sizes):
            prev_size = input_size if i == 0 else sizes[i - 1]
            layers.append(torch.nn.Linear(prev_size, size))
            layers.append(torch.nn.BatchNorm1d(size))
            layers.append(torch.nn.ReLU())
        layers.append(torch.nn.Linear(sizes[-1], output_size))
        layers.append(torch.nn.Dropout(p=0.2))
        layers.append(torch.nn.Sigmoid())
        self.model = torch.nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

model = MLP(1152, 294).to(device)
model.load_state_dict(torch.load('data/class-model.pth'))
model.eval()

with torch.no_grad():
    Y_pred = model(X).cpu().numpy()

with open('data/class-labels.json', 'r') as f:
    labels = json.load(f)
if Y_pred.shape[1] != len(labels):
    print('Selected label type doesn\'t match the model output')
    assert False

results = {
    'labels': labels,
    'bgcs': {},
}
for bgc_index, bgc_name in enumerate(X_names):
    y = Y_pred[bgc_index].tolist()
    label_indices = sorted(range(len(y)), key=lambda i: y[i], reverse=True)
    results['bgcs'][bgc_name] = label_indices[0:20] # top 20 label indices
print('Generated predictions for %d BGCs' % len(results['bgcs']))

with open(output_path, 'w') as f:
    json.dump(results, f)
print('Done')