FireKitCV/predict.py at master · fire717/FireKitCV · 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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
import os,argparse
import random

from firelib import init_all, FireModel, FireRunner, FireData

from config import cfg
import pandas as pd


def softmax(x):
    x_row_max = x.max(axis=-1)
    x_row_max = x_row_max.reshape(list(x.shape)[:-1]+[1])
    x = x - x_row_max
    x_exp = np.exp(x)
    x_exp_row_sum = x_exp.sum(axis=-1).reshape(list(x.shape)[:-1]+[1])
    softmax = x_exp / x_exp_row_sum
    return softmax


def predict(cfg):

    init_all(cfg)


    model = FireModel(cfg)


    data = FireData(cfg)
    # data.showTrainData()
    # b

    test_loader = data.getTestDataloader()


    runner = FireRunner(cfg, model)

    #print(model)
    runner.load_model(cfg['model_path'])


    res_dict = runner.predict(test_loader)
    print(len(res_dict))

    # to csv
    res_df = pd.DataFrame.from_dict(res_dict, orient='index', columns=['label'])
    res_df = res_df.reset_index().rename(columns={'index':'image_id'})
    res_df.to_csv(os.path.join(cfg['save_dir'], 'pre.csv'),
                                index=False,header=True)


def predictMerge(cfg):
    initFire(cfg)


    model = FireModel(cfg)


    data = FireData(cfg)
    # data.showTrainData()
    # b

    test_loader = data.getTestDataloader()
    runner1 = FireRunner(cfg, model)
    runner1.load_model('output/efficientnet-b6_e17_fold0_0.93368.pth')
    print("load model1, start running.")
    res_dict1 = runner1.predictRaw(test_loader)
    print(len(res_dict1))

    test_loader = data.getTestDataloader()
    runner2 = FireRunner(cfg, model)
    runner2.load_model('output/efficientnet-b6_e18_fold1_0.94537.pth')
    print("load model2, start running.")
    res_dict2 = runner2.predictRaw(test_loader)

    test_loader = data.getTestDataloader()
    runner3 = FireRunner(cfg, model)
    runner3.load_model('output/efficientnet-b6_e14_fold2_0.91967.pth')
    print("load model3, start running.")
    res_dict3 = runner3.predictRaw(test_loader)

    test_loader = data.getTestDataloader()
    runner4 = FireRunner(cfg, model)
    runner4.load_model('output/efficientnet-b6_e18_fold3_0.92239.pth')
    print("load model4, start running.")
    res_dict4 = runner4.predictRaw(test_loader)

    # test_loader = data.getTestDataloader()
    # runner5 = FireRunner(cfg, model)
    # runner5.load_model('output/efficientnet-b6_e17_fold0_0.93368.pth')
    # print("load model5, start running.")
    # res_dict5 = runner5.predictRaw(test_loader)


    res_dict = {}
    for k,v in res_dict1.items():
        #print(k,v)
        v1 =np.argmax(v+res_dict2[k]+res_dict3[k]+res_dict4[k])
        res_dict[k] = v1

    res_list = sorted(res_dict.items(), key = lambda kv: int(kv[0].split("_")[-1].split('.')[0]))
    print(len(res_list), res_list[0])

    # to csv
    # res_list_final = []
    # for res in res_list:
    #     res_list_final.append([res[0]]+res[1])
    # #res_df = pd.DataFrame.from_dict(res_dict, orient='index', columns=['type'])
    # #res_df = res_df.reset_index().rename(columns={'index':'id'})
    # res_df = DataFrame(res_list_final, columns=['id','type','color','toward'])


    # res_df.to_csv(os.path.join(cfg['save_dir'], 'result.csv'),
    #                             index=False,header=True)

    with open('result.csv', 'w', encoding='utf-8') as f:
        f.write('file,label\n')
        for i in range(len(res_list)):
            line = [res_list[i][0], str(res_list[i][1])]
            line = ','.join(line)
            f.write(line+"\n")


def predictTTA(cfg):

    pass


def predictMergeTTA(cfg):

    pass


def main(cfg):

    if cfg["merge"]:
        if cfg["TTA"]:
            predictMergeTTA(cfg)
        else:
            predictMerge(cfg)
    else:
        if cfg["TTA"]:
            predictTTA(cfg)
        else:
            predict(cfg)


if __name__ == '__main__':
    main(cfg)