TileDetection/inference.py at main · MySuperSoul/TileDetection · 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
from ai_hub import inferServer, log
import json
import torch
import mmdet
from mmdet.apis import init_detector, inference_detector
import torch.nn as nn
import numpy as np
import cv2
import time
import os

exp_compositions = {
    'model1': {
        'config':
        '/work/configs/tile_round2/cascade_s50_rfp_mstrain_augv2.py',
        'checkpoint':
        '/work/work_dirs/round2/swa_cascade_s50_rfp_mstrain_aug_v2/swa_model_12.pth'
    }
}
exp = exp_compositions['model1']
config_file = exp['config']
checkpoint_file = exp['checkpoint']


class MyServer(inferServer):

    def __init__(self, model, using_pair=True):
        super().__init__(model)
        log.i('Init myserver now')

        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        log.i('Current device: ', device)

        self.device = device
        self.using_pair = using_pair
        self.model = model.to(device)
        self.model.eval()

    def pre_process(self, request):
        #json process
        start_time = time.time()

        file = request.files['img']
        file_t = request.files['img_t']
        self.filename = file.filename
        file_data = file.read()

        img = cv2.imdecode(
            np.frombuffer(file_data, np.uint8), cv2.IMREAD_COLOR)

        end_time = time.time()
        print('Preprocess time: {} seconds'.format(end_time - start_time))

        return img

    def pridect(self, data):
        # remember here to fetch the first one
        start_time = time.time()

        predictions = inference_detector(model=self.model, img=data)[0]

        end_time = time.time()
        print('Inference time: {} seconds'.format(end_time - start_time))

        return predictions

    def post_process(self, data):
        predictions = data
        predict_rslt = []

        max_score = -1
        score_thr = 0.28

        for i, bboxes in enumerate(predictions):
            if len(bboxes) > 0 and i != 0:
                detect_label = i
                image_name = self.filename
                for bbox in bboxes:
                    xmin, ymin, xmax, ymax, score = bbox.tolist()
                    xmin, ymin, xmax, ymax = round(float(xmin), 2), round(
                        float(ymin), 2), round(float(xmax),
                                               2), round(float(ymax), 2)
                    max_score = max(max_score, score)

                    if xmax - xmin < 3 or ymax - ymin < 3:
                        continue

                    dict_instance = dict()
                    dict_instance['name'] = image_name
                    dict_instance['category'] = detect_label
                    dict_instance['score'] = round(float(score), 6)
                    dict_instance['bbox'] = [xmin, ymin, xmax, ymax]
                    predict_rslt.append(dict_instance)

        if max_score < score_thr:
            predict_rslt = []

        return predict_rslt

    def debug(self, filename):
        start_time = time.time()
        self.filename = filename
        img = cv2.imread(filename)

        end_time = time.time()
        print('Preprocess time: {} seconds'.format(end_time - start_time))
        data = self.pridect(img)
        data = self.post_process(data)
        return data


if __name__ == '__main__':
    model = init_detector(
        config=config_file, checkpoint=checkpoint_file, device='cpu')
    log.i('Init model success')
    myserver = MyServer(model=model, using_pair=False)
    myserver.run(debuge=False)