if args.cuda:
im_data = im_data.cuda()
im_info = im_info.cuda()
num_boxes = num_boxes.cuda()
gt_boxes = gt_boxes.cuda()
# make variable
im_data = Variable(im_data)
im_info = Variable(im_info)
num_boxes = Variable(num_boxes)
gt_boxes = Variable(gt_boxes)
# initilize the network here.
if args.net == 'vgg16':
fasterRCNN = vgg16(imdb.classes,
pretrained=True,
class_agnostic=args.class_agnostic)
else:
fasterRCNN = None
print("network is not defined")
pdb.set_trace()
fasterRCNN.create_architecture()
lr = cfg.TRAIN.LEARNING_RATE = args.lr
weight_decay = cfg.TRAIN.WEIGHT_DECAY = args.weight_decay
# tr_momentum = cfg.TRAIN.MOMENTUM
# tr_momentum = args.momentum
if args.cuda:
fasterRCNN.cuda()
raise Exception(
'There is no input directory for loading network from ' +
input_dir)
load_name = os.path.join(input_dir,
'faster_rcnn_{}_{}.pth'.format(150, 275))
print(load_name)
pascal_classes = np.asarray([
'__background__', 'airplane', 'ship', 'storage tank',
'baseball diamond', 'tennis court', 'basketball court',
'ground track field', 'harbor', 'bridge', 'vehicle'
])
# initilize the network here.
fasterRCNN = vgg16(pascal_classes,
pretrained=False,
class_agnostic=args.class_agnostic)
fasterRCNN.create_architecture()
print("load checkpoint %s" % (load_name))
if args.cuda > 0:
checkpoint = torch.load(load_name)
else:
checkpoint = torch.load(load_name,
map_location=(lambda storage, loc: storage))
fasterRCNN.load_state_dict(checkpoint['model'])
if 'pooling_mode' in checkpoint.keys():
cfg.POOLING_MODE = checkpoint['pooling_mode']
print('load model successfully!')
# initilize the tensor holder here.
im_info = im_info.cuda()
num_boxes = num_boxes.cuda()
gt_boxes = gt_boxes.cuda()
# make variable
im_data = Variable(im_data)
im_info = Variable(im_info)
num_boxes = Variable(num_boxes)
gt_boxes = Variable(gt_boxes)
if args.cuda:
cfg.CUDA = True
# initilize the network here.
if args.net == 'vgg16':
model = vgg16(imdb_classes, pretrained=True, pretrained_rfcn=True, class_agnostic=args.class_agnostic)
elif args.net == 'res101':
model = resnet(imdb_classes, 101, pretrained=True, pretrained_rfcn=True, class_agnostic=args.class_agnostic)
elif args.net == 'res50':
model = resnet(imdb_classes, 50, pretrained=True, pretrained_rfcn=True, class_agnostic=args.class_agnostic)
#print(model)
#input()
elif args.net == 'res152':
model = resnet(imdb_classes, 152, pretrained=True, pretrained_rfcn=True, class_agnostic=args.class_agnostic)
else:
print("network is not defined")
pdb.set_trace()
model.create_architecture()
draw = ImageDraw.Draw(img)
font = ImageFont.truetype("fonts/STFANGSO.TTF", 26, encoding="utf-8")
draw.text((bbox[0] + 5, bbox[1] + 20), txt, (249, 102, 183), font=font)
im = np.array(img)[:, :, ::-1]
else:
cv2.putText(im, txt, (bbox[0] + 5, bbox[1] + 20), cv2.FONT_HERSHEY_TRIPLEX,
0.5, (249, 102, 183), thickness=1)
return im
if __name__ == '__main__':
cfg_from_list(['ANCHOR_SCALES', '[8, 16, 32]',
'ANCHOR_RATIOS', '[0.5,1,2]'])
cfg.TRAIN.USE_FLIPPED = False
if net == 'vgg16':
fasterRCNN = vgg16(all_classes, pretrained=False,
class_agnostic=class_agnostic)
elif net == 'res101':
fasterRCNN = resnet(all_classes, 101, pretrained=False,
class_agnostic=class_agnostic)
elif net == 'res50':
fasterRCNN = resnet(all_classes, 50, pretrained=False,
class_agnostic=class_agnostic)
elif net == 'res152':
fasterRCNN = resnet(all_classes, 152, pretrained=False,
class_agnostic=class_agnostic)
else:
print("network is not defined")
pdb.set_trace()
fasterRCNN.create_architecture()
fasterRCNN.to(device)
def main():
args = parse_args()
print('Called with args:')
print(args)
best_meanap = 0
meanap = 0
if args.dataset == "virat":
args.set_cfgs = [
'ANCHOR_SCALES', '[1, 2, 3]', 'ANCHOR_RATIOS', '[0.5,1,2]',
'MAX_NUM_GT_BOXES', '15'
]
if args.dataset == "ucfsport":
args.set_cfgs = [
'ANCHOR_SCALES', '[4,8,16,24,28]', 'ANCHOR_RATIOS', '[0.5,1,2]',
'MAX_NUM_GT_BOXES', '2'
]
if args.dataset == "urfall":
args.set_cfgs = [
'ANCHOR_SCALES', '[4,8,16,24,28]', 'ANCHOR_RATIOS', '[0.5,1,2]',
'MAX_NUM_GT_BOXES', '1'
]
if args.dataset == "imfd":
args.set_cfgs = [
'ANCHOR_SCALES', '[4,8,16,24,28]', 'ANCHOR_RATIOS', '[0.5,1,2]',
'MAX_NUM_GT_BOXES', '1'
]
if args.dataset == "jhmdb":
args.set_cfgs = [
'ANCHOR_SCALES', '[4, 8,16,24,28]', 'ANCHOR_RATIOS', '[0.5,1,2]',
'MAX_NUM_GT_BOXES', '1'
]
if args.dataset == "ucf24":
args.set_cfgs = [
'ANCHOR_SCALES', '[4,8,16,24,28]', 'ANCHOR_RATIOS', '[0.5,1,2]',
'MAX_NUM_GT_BOXES', '1'
]
args.cfg_file = "cfgs/{}.yml".format(args.net)
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
np.random.seed(cfg.RNG_SEED)
#torch.backends.cudnn.benchmark = True
if torch.cuda.is_available() and not args.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
# train set
# -- Note: Use validation set and disable the flipped to enable faster loading.
cfg.USE_GPU_NMS = args.cuda
if args.dataset == 'virat':
num_class = cfg.VIRAT.NUM_CLASS
output_dir = cfg.VIRAT.output_model_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(output_dir):
os.makedirs(output_dir)
elif args.dataset == 'ucfsport':
num_class = cfg.UCFSPORT.NUM_CLASSES
output_dir = cfg.UCFSPORT.output_model_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(output_dir):
os.makedirs(output_dir)
elif args.dataset == 'urfall':
num_class = cfg.URFD.NUM_CLASSES
output_dir = cfg.URFD.output_model_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(output_dir):
os.makedirs(output_dir)
elif args.dataset == 'imfd':
num_class = cfg.IMFD.NUM_CLASSES
output_dir = cfg.IMFD.output_model_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(output_dir):
os.makedirs(output_dir)
elif args.dataset == 'jhmdb':
num_class = cfg.JHMDB.NUM_CLASSES
output_dir = cfg.JHMDB.output_model_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(output_dir):
os.makedirs(output_dir)
elif args.dataset == 'ucf24':
num_class = cfg.UCF24.NUM_CLASSES
output_dir = cfg.UCF24.output_model_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(output_dir):
os.makedirs(output_dir)
else:
print("dataset is not defined ")
#log initialisation
args.store_name = '_'.join([
'S-RAD', args.dataset, args.net,
'segment%d' % args.num_segments, 'e{}'.format(args.max_epochs),
'session%d' % args.session
])
check_rootfolders(args.store_name, args.dataset)
#logging
log_training, logger = log_info(cfg,
args.store_name,
args.dataset,
args=args)
#dataloader
train_loader = construct_loader(cfg,
dataset=args.dataset,
num_segments=args.num_segments,
interval=args.interval,
batch_size=args.batch_size,
split='train',
input_sampling=True,
split_num=args.splits,
pathway=args.pathway)
val_loader = construct_loader(cfg,
dataset=args.dataset,
num_segments=args.num_segments,
interval=args.interval,
batch_size=args.batch_size,
split='val',
input_sampling=True,
split_num=args.splits,
pathway=args.pathway)
if args.dataset == 'virat':
test_loader = construct_loader(cfg,
dataset=args.dataset,
num_segments=args.num_segments,
interval=args.interval,
batch_size=args.batch_size,
split='test',
input_sampling=True,
split_num=args.splits,
pathway=args.pathway)
# prevent something not % n_GPU
if args.cuda:
cfg.CUDA = True
# initilize the network here.
if args.net == 'vgg16':
S_RAD = vgg16(num_class,
pretrained=True,
class_agnostic=args.class_agnostic,
loss_type=args.loss_type)
elif args.net == 'res50':
S_RAD = resnet(num_class,
num_layers=50,
base_model='resnet50',
n_segments=args.num_segments,
n_div=args.shift_div,
place=args.shift_place,
pretrain=args.pretrain,
shift=args.shift,
class_agnostic=args.class_agnostic,
loss_type=args.loss_type,
pathway=args.pathway)
else:
print("network is not defined")
pdb.set_trace()
#create the architecture
S_RAD.create_architecture()
#set the parameters
lr = args.lr
params = []
for key, value in dict(S_RAD.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
params += [{'params':[value],'lr':lr*(cfg.TRAIN.DOUBLE_BIAS + 1), \
'weight_decay': cfg.TRAIN.BIAS_DECAY and cfg.TRAIN.WEIGHT_DECAY or 0}]
else:
params += [{
'params': [value],
'lr': lr,
'weight_decay': cfg.TRAIN.WEIGHT_DECAY
}]
if args.cuda:
S_RAD.cuda()
#define optimizer
if args.optimizer == "adam":
lr = lr * 0.1
optimizer = torch.optim.Adam(params)
elif args.optimizer == "sgd":
optimizer = torch.optim.SGD(params, momentum=cfg.TRAIN.MOMENTUM)
#adding UCF sport weights to the first branch base1
#if args.pathway == "two_pathway":
#if args.tune_from:
## adding temporal shift pretrained kinetics weights
#if args.pathway =="naive":
if args.tune_from:
print(("=> fine-tuning from '{}'".format(args.tune_from)))
sd = torch.load(args.tune_from)
sd = sd['state_dict']
model_dict = S_RAD.state_dict()
replace_dict = []
for k, v in sd.items():
if k not in model_dict:
replace_dict.append(
(k.replace(
'module.base_model.conv1', 'RCNN_base1.0').replace(
'module.base_model.bn1', 'RCNN_base1.1').replace(
'module.base_model.layer1.0',
'RCNN_base1.4.0').replace(
'module.base_model.layer1.1',
'RCNN_base1.4.1').replace(
'module.base_model.layer1.2',
'RCNN_base1.4.2').replace(
'module.base_model.layer2.0',
'RCNN_base1.5.0').
replace('module.base_model.layer2.1',
'RCNN_base1.5.1').replace(
'module.base_model.layer2.2',
'RCNN_base1.5.2').replace(
'module.base_model.layer2.3',
'RCNN_base1.5.3').replace(
'module.base_model.layer3.0',
'RCNN_base1.6.0').replace(
'module.base_model.layer3.1',
'RCNN_base1.6.1').replace(
'module.base_model.layer3.2',
'RCNN_base1.6.2').
replace('module.base_model.layer3.3',
'RCNN_base1.6.3').replace(
'module.base_model.layer3.4',
'RCNN_base1.6.4').replace(
'module.base_model.layer3.5',
'RCNN_base1.6.5').replace(
'module.base_model.layer4.0.',
'RCNN_top.0.0.').replace(
'module.base_model.layer4.1',
'RCNN_top.0.1').replace(
'module.base_model.layer4.2',
'RCNN_top.0.2').
replace('module.base_model.layer4.0.conv1.net',
'RCNN_top.0.0.conv1').replace(
'module.base_model.layer4.1.conv1.net',
'RCNN_top.0.1.conv1').replace(
'module.base_model.layer4.2.conv1.net',
'RCNN_top.0.2.conv1').replace(
'RCNN_top.0.0.conv1.net',
'RCNN_top.0.0.conv1').replace(
'RCNN_top.0.1.conv1.net',
'RCNN_top.0.1.conv1').replace(
'RCNN_top.0.2.conv1.net',
'RCNN_top.0.2.conv1'), k))
for k_new, k in replace_dict:
sd[k_new] = sd.pop(k)
keys1 = set(list(sd.keys()))
keys2 = set(list(model_dict.keys()))
set_diff = (keys1 - keys2) | (keys2 - keys1)
print('#### Notice: keys that failed to load: {}'.format(set_diff))
if args.dataset not in args.tune_from: # new dataset
print('=> New dataset, do not load fc weights')
sd = {k: v for k, v in sd.items() if 'fc' not in k}
model_dict.update(sd)
S_RAD.load_state_dict(model_dict)
if args.resume:
load_name = os.path.join(
output_dir,
'S-RAD_{}_{}_{}.pth'.format(args.checksession, args.checkepoch,
args.checkpoint))
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name)
args.session = checkpoint['session']
args.start_epoch = checkpoint['epoch']
sd = checkpoint['model']
model_dict = S_RAD.state_dict()
replace_dict = []
for k, v in sd.items():
if k not in model_dict:
replace_dict.append((k.replace('RCNN_base', 'RCNN_base1'), k))
for k_new, k in replace_dict:
sd[k_new] = sd.pop(k)
keys1 = set(list(sd.keys()))
keys2 = set(list(model_dict.keys()))
set_diff = (keys1 - keys2) | (keys2 - keys1)
print('#### Notice: keys that failed to load: {}'.format(set_diff))
model_dict.update(sd)
S_RAD.load_state_dict(model_dict)
optimizer.load_state_dict(checkpoint['optimizer'])
lr = optimizer.param_groups[0]['lr']
if 'pooling_mode' in checkpoint.keys():
cfg.POOLING_MODE = checkpoint['pooling_mode']
print("loaded checkpoint %s" % (load_name))
if args.mGPUs:
mGPUs = True
S_RAD = nn.DataParallel(S_RAD)
else:
mGPUs = False
session = args.session
if args.evaluate:
validate_voc(val_loader, S_RAD,args.start_epoch,num_class, \
args.num_segments,session,args.batch_size,\
cfg,log_training,args.dataset,args.pathway,args.eval_metrics)
sys.exit
for epoch in range(args.start_epoch, args.max_epochs + 1):
if epoch % (args.lr_decay_step + 1) == 0:
adjust_learning_rate(optimizer, args.lr_decay_gamma)
lr *= args.lr_decay_gamma
if args.dataset == 'virat':
#dataloader
train(train_loader, S_RAD, lr, optimizer, epoch, num_class,
args.batch_size, session, mGPUs, logger, output_dir, data,
cfg, args.acc_step, log_training)
# evaluate on validation set
validate_virat(val_loader, S_RAD,epoch,num_class, \
args.num_segments,session,args.batch_size,data,cfg,\
log_training,args.dataset)
if epoch % 10 == 0:
validate_virat(test_loader, S_RAD,epoch,num_class, \
args.num_segments,session,args.batch_size,data,cfg,
log_training,args.dataset)
elif args.dataset == 'ucfsport':
train(train_loader, S_RAD, lr, optimizer, epoch, num_class,
args.batch_size, session, mGPUs, logger, output_dir, cfg,
args.acc_step, log_training)
validate_voc(val_loader, S_RAD,epoch,num_class, \
args.num_segments,session,args.batch_size,
cfg,log_training,args.dataset,args.pathway,args.eval_metrics)
elif args.dataset == 'urfall':
train(train_loader, S_RAD, lr, optimizer, epoch, num_class,
args.batch_size, session, mGPUs, logger, output_dir, cfg,
args.acc_step, log_training)
validate_voc(val_loader, S_RAD,epoch,num_class, \
args.num_segments,session,args.batch_size,\
cfg,log_training,args.dataset,args.pathway,args.eval_metrics)
elif args.dataset == 'imfd':
train(train_loader, S_RAD, lr, optimizer, epoch, num_class,
args.batch_size, session, mGPUs, logger, output_dir, cfg,
args.acc_step, log_training)
validate_voc(val_loader, S_RAD,epoch,num_class, \
args.num_segments,session,args.batch_size,
cfg,log_training,args.dataset,args.pathway,args.eval_metrics)
elif args.dataset == 'jhmdb':
train(train_loader, S_RAD, lr, optimizer, epoch, num_class,
args.batch_size, session, mGPUs, logger, output_dir, cfg,
args.acc_step, log_training)
if epoch % 2 == 0:
validate_voc(val_loader, S_RAD,epoch,num_class, \
args.num_segments,session,args.batch_size,
cfg,log_training,args.dataset,args.pathway,args.eval_metrics)
elif args.dataset == 'ucf24':
train(train_loader, S_RAD, lr, optimizer, epoch, num_class,
args.batch_size, session, mGPUs, logger, output_dir, cfg,
args.acc_step, log_training)
validate_voc(val_loader, S_RAD,epoch,num_class, \
args.num_segments,session,args.batch_size,
cfg,log_training,args.dataset,args.pathway,args.eval_metrics)
cfg_from_list(args.set_cfgs)
# print('Config for faster-rcnn:')
# pprint.pprint(cfg)
''' Initialize FasterRCNN model '''
if args.net == 'vgg16' or args.net == 'res101':
print('Initialize fasterRCNN module')
else:
print('FasterRCNN model is not defined')
pdb.set_trace()
# 20 object classes in PascalVOC (plus '__background__')
pascal_category = np.asarray(['__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person',
'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'])
if args.feature_extraction_cnn == 'vgg16':
FasterRCNN = vgg16(pascal_category, use_pytorch=False, use_caffe=False, class_agnostic=args.class_agnostic)
elif args.feature_extraction_cnn == 'resnet101':
FasterRCNN = resnet(pascal_category, 101, use_pytorch=False, use_caffe=False, class_agnostic=args.class_agnostic)
FasterRCNN.create_architecture()
''' Set FasterRCNN module '''
# The directory for loading pre-trained fasterRCNN for extracting bounding box of objects
RCNN_cp_dir = os.path.join(args.load_dir, args.net, args.dataset)
RCNN_cp_name = os.path.join(RCNN_cp_dir, 'best_faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint))
print('Load fasterRCNN model {}'.format(RCNN_cp_name))
if not os.path.exists(RCNN_cp_name):
raise Exception('There is no pre-trained fasterRCNN model, i.e. {}'.format(RCNN_cp_name))
# Load pre-trained parameters for fasterRCNN
RCNN_checkpoint = torch.load(RCNN_cp_name, map_location=lambda storage, loc: storage)
# for name, param in model.FasterRCNN.state_dict().items():
# model.FasterRCNN.state_dict()[name].copy_(RCNN_checkpoint['model'][name])
def evaluation(name, net=None, vis=False, cuda=True, class_agnostic=False):
cfg.TRAIN.USE_FLIPPED = False
imdb, roidb, ratio_list, ratio_index = combined_roidb(name, False)
imdb.competition_mode(on=True)
print('{:d} roidb entries'.format(len(roidb)))
if not net:
input_dir = args.load_dir + "/" + args.net + "/" + args.dataset
# input_dir = 'weight'
if not os.path.exists(input_dir):
raise Exception(
'There is no input directory for loading network from ' +
input_dir)
# load_name = os.path.join(input_dir,
# 'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint))
load_name = os.path.join(
input_dir, 'faster_rcnn_{}_best.pth'.format(cfg['POOLING_MODE']))
# initilize the network here.
if args.net == 'vgg16':
fasterRCNN = vgg16(imdb.classes,
pretrained=False,
class_agnostic=args.class_agnostic)
elif args.net == 'res101':
fasterRCNN = resnet(imdb.classes,
101,
pretrained=False,
class_agnostic=args.class_agnostic)
elif args.net == 'res50':
fasterRCNN = resnet(imdb.classes,
50,
pretrained=False,
class_agnostic=args.class_agnostic)
elif args.net == 'res152':
fasterRCNN = resnet(imdb.classes,
152,
pretrained=False,
class_agnostic=args.class_agnostic)
else:
print("network is not defined")
pdb.set_trace()
fasterRCNN.create_architecture()
print("load checkpoint %s" % (load_name))
checkpoint = torch.load(load_name)
fasterRCNN.load_state_dict(checkpoint['model'])
if 'pooling_mode' in checkpoint.keys():
cfg.POOLING_MODE = checkpoint['pooling_mode']
print('load model successfully!')
else:
fasterRCNN = net
# initilize the tensor holder here.
im_data = torch.FloatTensor(1)
im_info = torch.FloatTensor(1)
num_boxes = torch.LongTensor(1)
gt_boxes = torch.FloatTensor(1)
# ship to cuda
if cuda:
im_data = im_data.cuda()
im_info = im_info.cuda()
num_boxes = num_boxes.cuda()
gt_boxes = gt_boxes.cuda()
# make variable
im_data = Variable(im_data)
im_info = Variable(im_info)
num_boxes = Variable(num_boxes)
gt_boxes = Variable(gt_boxes)
if cuda:
cfg.CUDA = True
if cuda:
fasterRCNN.cuda()
start = time.time()
max_per_image = 100
# vis = args.vis
if vis:
thresh = 0.05
else:
thresh = 0.0
save_name = 'faster_rcnn_10'
num_images = len(imdb.image_index)
all_boxes = [[[] for _ in range(num_images)]
for _ in range(imdb.num_classes)]
output_dir = get_output_dir(imdb, save_name)
dataset = roibatchLoader(roidb, ratio_list, ratio_index, 1, \
imdb.num_classes, training=False, normalize=False)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=0,
pin_memory=True)
data_iter = iter(dataloader)
_t = {'im_detect': time.time(), 'misc': time.time()}
det_file = os.path.join(output_dir, 'detections.pkl')
fasterRCNN.eval()
empty_array = np.transpose(np.array([[], [], [], [], []]), (1, 0))
for i in range(num_images):
data = next(data_iter)
with torch.no_grad():
im_data.resize_(data[0].size()).copy_(data[0])
im_info.resize_(data[1].size()).copy_(data[1])
gt_boxes.resize_(data[2].size()).copy_(data[2])
num_boxes.resize_(data[3].size()).copy_(data[3])
det_tic = time.time()
rois, cls_prob, bbox_pred, \
rpn_loss_cls, rpn_loss_box, \
RCNN_loss_cls, RCNN_loss_bbox, \
rois_label = fasterRCNN(im_data, im_info, gt_boxes, num_boxes)
scores = cls_prob.data
boxes = rois.data[:, :, 1:5]
if cfg.TEST.BBOX_REG:
# Apply bounding-box regression deltas
box_deltas = bbox_pred.data
if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED:
# Optionally normalize targets by a precomputed mean and stdev
if class_agnostic:
box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
+ torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
box_deltas = box_deltas.view(1, -1, 4)
else:
box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \
+ torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
box_deltas = box_deltas.view(1, -1, 4 * len(imdb.classes))
pred_boxes = bbox_transform_inv(boxes, box_deltas, 1)
pred_boxes = clip_boxes(pred_boxes, im_info.data, 1)
else:
# Simply repeat the boxes, once for each class
pred_boxes = np.tile(boxes, (1, scores.shape[1]))
pred_boxes /= data[1][0][2].item()
scores = scores.squeeze()
pred_boxes = pred_boxes.squeeze()
det_toc = time.time()
detect_time = det_toc - det_tic
misc_tic = time.time()
if vis:
im = cv2.imread(imdb.image_path_at(i))
im2show = np.copy(im)
for j in range(1, imdb.num_classes):
inds = torch.nonzero(scores[:, j] > thresh).view(-1)
# if there is det
if inds.numel() > 0:
cls_scores = scores[:, j][inds]
_, order = torch.sort(cls_scores, 0, True)
if class_agnostic:
cls_boxes = pred_boxes[inds, :]
else:
cls_boxes = pred_boxes[inds][:, j * 4:(j + 1) * 4]
cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1)
# cls_dets = torch.cat((cls_boxes, cls_scores), 1)
cls_dets = cls_dets[order]
keep = nms(cls_boxes[order, :], cls_scores[order],
cfg.TEST.NMS)
cls_dets = cls_dets[keep.view(-1).long()]
if vis:
im2show = vis_detections(im2show, imdb.classes[j],
cls_dets.cpu().numpy(), 0.3)
all_boxes[j][i] = cls_dets.cpu().numpy()
else:
all_boxes[j][i] = empty_array
# Limit to max_per_image detections *over all classes*
if max_per_image > 0:
image_scores = np.hstack(
[all_boxes[j][i][:, -1] for j in range(1, imdb.num_classes)])
if len(image_scores) > max_per_image:
image_thresh = np.sort(image_scores)[-max_per_image]
for j in range(1, imdb.num_classes):
keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0]
all_boxes[j][i] = all_boxes[j][i][keep, :]
misc_toc = time.time()
nms_time = misc_toc - misc_tic
sys.stdout.write('im_detect: {:d}/{:d} {:.3f}s {:.3f}s \r' \
.format(i + 1, num_images, detect_time, nms_time))
sys.stdout.flush()
if vis:
cv2.imwrite('result.png', im2show)
pdb.set_trace()
# cv2.imshow('test', im2show)
# cv2.waitKey(0)
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
map = imdb.evaluate_detections(all_boxes, output_dir)
# print(map)
end = time.time()
print("test time: %0.4fs" % (end - start))
return map
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
args.cfg_file = "cfgs/{}_ls.yml".format(args.net) if args.large_scale else "cfgs/{}.yml".format(args.net)
args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]', 'MAX_NUM_GT_BOXES', '20']
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
print(cfg['POOLING_MODE'])
# Set up model
classes = ('__background__', # always index 0
'crazing', 'inclusion', 'patches',
'pitted_surface', 'rolled-in_scale', 'scratches')
fasterRCNN = vgg16(classes, pretrained=False, class_agnostic=False)
fasterRCNN.create_architecture()
fasterRCNN.cuda()
fasterRCNN.eval()
im_data = torch.FloatTensor(1).to(device)
im_info = torch.FloatTensor(1).to(device)
num_boxes = torch.LongTensor(1).to(device)
gt_boxes = torch.FloatTensor(1).to(device)
input = np.random.random((1,args.img_size, args.img_size,3))
im_blob = input
im_info_np = np.array([[im_blob.shape[1], im_blob.shape[2], 1.5]], dtype=np.float32)
im_data_pt = torch.from_numpy(im_blob)
im_data_pt = im_data_pt.permute(0, 3, 1, 2)
print('{:d} roidb entries'.format(len(roidb)))
input_dir = args.load_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(input_dir):
raise Exception(
'There is no input directory for loading network from ' +
input_dir)
load_name = os.path.join(
input_dir,
'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch,
args.checkpoint))
# initilize the network here.
if args.net == 'vgg16':
fasterRCNN = vgg16(num_class,
pretrained=False,
class_agnostic=args.class_agnostic)
elif args.net == 'res50':
fasterRCNN = resnet(num_class,
50,
pretrained=False,
class_agnostic=args.class_agnostic)
elif args.net == 'res101':
fasterRCNN = resnet(num_class,
101,
pretrained=True,
class_agnostic=args.class_agnostic)
elif args.net == 'res152':
fasterRCNN = resnet(num_class,
152,
pretrained=True,
def load_detector(dataset):
args = parse_args(dataset)
print('Called with args:')
print(args)
np.random.seed(cfg.RNG_SEED)
if args.dataset == "vg":
args.imdb_name = "vg_2007_trainval"
args.imdbval_name = "vg_2007_trainval"
args.set_cfgs = [
'ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'
]
from open_relation.dataset.vg.label_hier.obj_hier import objnet
classes = np.asarray(objnet.get_raw_labels())
elif args.dataset == "vrd":
args.imdb_name = "vrd_2007_trainval"
args.imdbval_name = "vrd_2007_test"
args.set_cfgs = [
'ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'
]
from open_relation.dataset.vrd.label_hier.obj_hier import objnet
classes = np.asarray(objnet.get_raw_labels())
elif args.dataset == "pascal_voc":
args.imdb_name = "voc_2007_trainval"
args.imdbval_name = "voc_2007_test"
args.set_cfgs = [
'ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'
]
classes = np.asarray([
'__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train',
'tvmonitor'
])
args.cfg_file = os.path.join(
project_root, "cfgs/{}_ls.yml".format(args.net)
if args.large_scale else "cfgs/{}.yml".format(args.net))
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
print('Using config:')
pprint.pprint(cfg)
input_dir = args.load_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(input_dir):
raise Exception(
'There is no input directory for loading network from ' +
input_dir)
load_name = os.path.join(
input_dir,
'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch,
args.checkpoint))
# initilize the network here.
fasterRCNN = vgg16(classes,
pretrained=False,
class_agnostic=args.class_agnostic)
fasterRCNN.create_architecture()
print("load checkpoint %s" % (load_name))
checkpoint = torch.load(load_name)
fasterRCNN.load_state_dict(checkpoint['model'])
if 'pooling_mode' in checkpoint.keys():
cfg.POOLING_MODE = checkpoint['pooling_mode']
print('load model successfully!')
if args.cuda:
cfg.CUDA = True
if args.cuda:
fasterRCNN.cuda()
fasterRCNN.eval()
return fasterRCNN
