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
# from tensorrtserver.api import *
import os
import sys
import argparse
import torch
import horovod.torch as hvd
import numpy as np
import lmdb
import json
import SimpleITK as sitk
import time
if __name__ == "__main__" and __package__ is None:
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".."))
__package__ = "lungDetection3D"
from lungDetection3D.BaseDetector.config import get_cfg_defaults
from lungDetection3D.BaseDetector.utils.gpu_utils import set_gpu
from lungDetection3D.NoduleDetector.engine import NoduleDetection3D
from lungDetection3D.NoduleDetector.modeling.detector.nodulenet import NoduleNet
from lungDetection3D.NoduleDetector.data.bbox_reader import BboxReader_Nodule
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="full functional execute script of Detection module.")
group = parser.add_mutually_exclusive_group()
group.add_argument("-t", "--train", help="model training", action="store_true")
group.add_argument("-i", "--inference", help="model inference", action="store_true")
parser.add_argument("-c", "--config", type=str, default="config.yaml", help="config file path")
args = parser.parse_args()
cfg = get_cfg_defaults()
cfg.merge_from_file(args.config)
cfg.freeze()
# load model
os.environ['CUDA_VISIBLE_DEVICES'] = '7'
# checkpoint = "/data/huangky/FPN/20201119/checkpoints/069.pth"
# model = NoduleNet(cfg, mode='test')
# model.eval()
# device = torch.device("cuda")
# model.to(device)
# checkpoint = torch.load(checkpoint)
# model.load_state_dict(checkpoint['state_dict'])
# print("model successfully loaded")
model_path = '/data/huangky/FPN/20201119/NoduleNet_FPN_half.pt'
model = torch.jit.load(model_path)
# generate data
test_set_name = "hahaha"
dataset = BboxReader_Nodule(cfg, mode='test')
for i, (input, image) in enumerate(dataset):
D, H, W = image.shape
pid = dataset.sample_bboxes[i].get_field("filename")
pid = pid.split('/')[-2].replace('.nii.gz', '')
print('[%d] Predicting %s' % (i, pid), image.shape)
input = input.cuda().unsqueeze(0).half()
# rpns, detections, ensembles = model.forward(input, [])
# torchscript infer
# input_data = torch.from_numpy(input_data).cuda()
# rpn_windows = torch.from_numpy(rpn_windows).cuda()
# model = torch.jit.load(model_path)
with torch.no_grad():
rpns = model(input)
time.sleep(5)
if cfg.MODEL.BACKBONE.FPN:
fpns = rpns
for i in range(len(fpns)):
fpn = fpns[i]
fpn = fpn.cpu().numpy()
if len(fpn):
fpn = fpn[:, 1:]
print("-------------------torchscript infer-------------------")
print('fpns_%s' % str(i))
print("fpns: {}".format(fpn))
torch.cuda.empty_cache()
# # pth to torchscript
# model_path = '/data/huangky/FPN/20201119/NoduleNet_FPN.pt'
# with torch.no_grad():
# traced_model = torch.jit.trace(model, (input))
# print(traced_model.graph)
# traced_model.save(model_path)