Spaces:
Running
Running
File size: 3,204 Bytes
5263bd3 a61f4be 870813f 5263bd3 870813f 2243c0c 870813f a61f4be 870813f 63d967d 870813f 63d967d 870813f 5263bd3 870813f 2243c0c 870813f 5263bd3 870813f 5263bd3 870813f 5263bd3 870813f 5263bd3 870813f 5263bd3 63d967d |
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 |
import gradio as gr
import torch
import joblib
import numpy as np
from itertools import product
import torch.nn as nn
class VirusClassifier(nn.Module):
def __init__(self, input_shape: int):
super(VirusClassifier, self).__init__()
self.network = nn.Sequential(
nn.Linear(input_shape, 64),
nn.GELU(),
nn.BatchNorm1d(64),
nn.Dropout(0.3),
nn.Linear(64, 32),
nn.GELU(),
nn.BatchNorm1d(32),
nn.Dropout(0.3),
nn.Linear(32, 32),
nn.GELU(),
nn.Linear(32, 2)
)
def forward(self, x):
return self.network(x)
def sequence_to_kmer_vector(sequence: str, k: int = 6) -> np.ndarray:
kmers = [''.join(p) for p in product("ACGT", repeat=k)]
kmer_dict = {kmer: 0 for kmer in kmers}
for i in range(len(sequence) - k + 1):
kmer = sequence[i:i+k]
if kmer in kmer_dict:
kmer_dict[kmer] += 1
return np.array(list(kmer_dict.values()))
def parse_fasta(text):
sequences = []
current_header = None
current_sequence = []
for line in text.split('\n'):
line = line.strip()
if not line:
continue
if line.startswith('>'):
if current_header:
sequences.append((current_header, ''.join(current_sequence)))
current_header = line[1:]
current_sequence = []
else:
current_sequence.append(line.upper())
if current_header:
sequences.append((current_header, ''.join(current_sequence)))
return sequences
def predict(file_obj):
if file_obj is None:
return "Please upload a FASTA file"
# Read the file content
text = file_obj.read().decode()
# Load model and scaler
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = VirusClassifier(4096).to(device)
model.load_state_dict(torch.load('model.pt', map_location=device))
scaler = joblib.load('scaler.pkl')
model.eval()
# Get predictions
results = []
sequences = parse_fasta(text)
for header, seq in sequences:
# Get k-mer vector
kmer_vector = sequence_to_kmer_vector(seq)
kmer_vector = scaler.transform(kmer_vector.reshape(1, -1))
# Predict
with torch.no_grad():
output = model(torch.FloatTensor(kmer_vector).to(device))
probs = torch.softmax(output, dim=1)
# Format results
pred_class = 1 if probs[0][1] > probs[0][0] else 0
pred_label = 'human' if pred_class == 1 else 'non-human'
result = f"""
Sequence: {header}
Prediction: {pred_label}
Confidence: {float(max(probs[0])):0.4f}
Human probability: {float(probs[0][1]):0.4f}
Non-human probability: {float(probs[0][0]):0.4f}
"""
results.append(result)
return "\n".join(results)
# Create the interface
iface = gr.Interface(
fn=predict,
inputs=gr.File(label="Upload FASTA file"),
outputs=gr.Textbox(label="Results"),
title="Virus Host Classifier"
)
# Launch with public link
iface.launch(share=True) |