def __init__(self, args):
torch.manual_seed(317)
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
heads = {'hm': args.num_classes,
'reg': 2*args.num_classes,
'wh': 2*4,}
self.model = spinal_net.SpineNet(heads=heads,
pretrained=True,
down_ratio=args.down_ratio,
final_kernel=1,
head_conv=256)
self.num_classes = args.num_classes
self.decoder = decoder.DecDecoder(K=args.K, conf_thresh=args.conf_thresh)
self.dataset = {'spinal': BaseDataset}
def __init__(self):
torch.manual_seed(317)
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
heads = {
'hm': 1,
'reg': 2,
'wh': 2 * 4,
}
self.model_path = "model_last.pth"
self.model = spinal_net.SpineNet(heads=heads,
basename='resnet34',
pretrained=True,
down_ratio=4,
final_kernel=1,
head_conv=256)
self.down_ratio = 4
self.model_state = False
self.decoder = decoder.DecDecoder(K=100, conf_thresh=0.2)
num_classes = {'dota': 15, 'hrsc': 1}
heads = {
'hm': num_classes[args.dataset],
'wh': 10,
'reg': 2,
'cls_theta': 1
}
down_ratio = 4
model = ctrbox_net.CTRBOX(heads=heads,
pretrained=True,
down_ratio=down_ratio,
final_kernel=1,
head_conv=256)
decoder = decoder.DecDecoder(K=args.K,
conf_thresh=args.conf_thresh,
num_classes=num_classes[args.dataset])
if args.phase == 'train':
ctrbox_obj = train.TrainModule(dataset=dataset,
num_classes=num_classes,
model=model,
decoder=decoder,
down_ratio=down_ratio)
ctrbox_obj.train_network(args)
elif args.phase == 'test':
ctrbox_obj = test.TestModule(dataset=dataset,
num_classes=num_classes,
model=model,
decoder=decoder)
ctrbox_obj.test(args, down_ratio=down_ratio)
cfg = Config()
mkdir(cfg.out_dir)
txtPath = os.path.join(cfg.data_root, 'valid.txt')
imgList = []
f = open(txtPath)
datas = f.readlines()
for data in datas:
path = data.split(' ')[0]
imgList.append(path)
imgList = [os.path.join(cfg.data_root, 'images', f'{imgName}.jpg') for imgName in imgList]
a = imgList.sort()
net = load_net(cfg).to(device)
net.eval()
decoder = decoder.DecDecoder(K=cfg.K,
conf_thresh=cfg.conf_thresh,
num_classes=cfg.num_classes)
frame_size = cfg.image_size - cfg.gap
for j, imgPath in tqdm.tqdm(enumerate(imgList)):
image_name = os.path.split(imgPath)[-1].split('.')[0]
image = cv2.imread(imgPath, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
img2, ratio, pad = letterbox(image.copy(), (cfg.image_size, cfg.image_size), auto=False, scaleup=False)
sample = img2.copy()
img2 = img2.astype(np.float32) / 255.
img2 -= 0.5
img2 = img2.transpose(2, 0, 1).reshape(1, 3, cfg.image_size, cfg.image_size)
img2 = torch.from_numpy(img2).to(device)
bboxes, scores = make_predictions(net, img2, cfg)
picked_boxes, picked_score = nms(bboxes, scores, threshold=0.3)