dnai/predict.py at main · alexsteinerr/dnai · GitHub

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
116
117
118
119
120
121
122
123
124
125
126
import torch
import models
import generate_pdf
import generate
import torch.nn.functional as F
import numpy as np

device = torch.device('cuda:'+ str(0) if torch.cuda.is_available() else 'cpu')

if torch.cuda.is_available():
    CNN = models.CNN().cuda()
else:
    CNN = models.CNN().cpu()

cnn_path = "models.pt"

checkpoint = torch.load(cnn_path,  map_location=device)
CNN.load_state_dict(checkpoint['CNN'])

length = 32768

def convertToArray(x):
    sequence = []

    '''
    Possible combination
        - AT
        - TA
        - CG
        - GC
    '''

    for i in range(len(x)):
        if i < len(x) - 1 and i % 2 == 0:
            if x[i] + x[i+1] == 'AT':
                sequence.append(0)
            elif x[i] + x[i+1] == 'TA':
                sequence.append(1)
            elif x[i] + x[i+1] == "CG":
                sequence.append(2)
            else:
                sequence.append(3)

    return np.array(sequence, dtype=np.float32)

def convertToString(x):
    '''
    Possible combination
        - AT
        - TA
        - CG
        - GC
    '''

    if x == 0:
        return "AT"
    elif x == 1:
        return "TA"
    elif x == 2:
        return "CG"
    else:
        return "GC"

def predict(sequence, count):
    print("Predicting...")

    sequence = sequence.upper()

    x = torch.from_numpy(convertToArray(sequence)).repeat(32,1,1).reshape(32,length,1).to(device)

    results = []
    patterns = f"<p class=MsoNormal style='text-align:justify><span lang=EN-GB style='font-size:1.0pt line-height:115%'></span>"

    ouput, activation_maps = CNN(x)

    print(ouput[0])

    print("Prediction Done!")

    sequence = convertToArray(sequence)

    for i, j in enumerate(sequence):
        if torch.argmax(ouput) != 0:
            if sequence[i] == activation_maps[0][1][i] == True and (ouput[0][1].item() > 10 or ouput[0][1].item() == torch.argmax(ouput)):
                print(1)
                patterns = patterns + f"<mark style='background-color: yellow;'>{convertToString(sequence[i])}</mark>"

            elif sequence[i] == activation_maps[0][2][i] == True and (ouput[0][2].item() > 10 or ouput[0][2].item() == torch.argmax(ouput)):
                print(2)
                patterns = patterns + f"<mark style='background-color: lightgreen;'>{convertToString(sequence[i])}</mark>"

            elif sequence[i] == activation_maps[0][3][i] == True and (ouput[0][3].item() > 10 or ouput[0][3].item() == torch.argmax(ouput)):
                print(3)
                patterns = patterns + f"<mark style='background-color: lightblue;'>{convertToString(sequence[i])}</mark>"

            else:
                print(4)
                patterns = patterns + f"{convertToString(sequence[i])}"

        else:
            print(0)
            if sequence[i] == activation_maps[0][0][i] == True and (ouput[0][0].item() > 10 or ouput[0][0].item() == torch.argmax(ouput)):
                patterns = patterns + f"<mark style='background-color: orange;'>{convertToString(sequence[i])}</mark>"

            else:
                patterns = patterns + f"{convertToString(sequence[i])}"

    if ouput[0][1].item() > 10 or ouput[0][1].item() == torch.argmax(ouput):
        results.append("lactose intolerance")
        print(1)
    if ouput[0][2].item() > 10 or ouput[0][2].item() == torch.argmax(ouput):
        results.append("haemophilia")
    if ouput[0][3].item() > 10 or ouput[0][3].item() == torch.argmax(ouput):
        results.append("autism")
    if ouput[0][1].item() == torch.argmax(ouput):
        results = ["no diseases found in the provided DNA sequence."]

    patterns = patterns + "<o:p></o:p></span></p>"

    print(results)

    generate_pdf.generate_pdf(results, patterns, count)

    return results

#predict(generate.generate_random_dna_sequence(length),1)