def create_npjson(s, image, json_path,use_gpu = True):
bboxes = list()
model=s.model
preproc= s.test_loader.dataset.preproc
detector = s.detector
model.eval()
# model.onnx_export = True
num_classes = detector.num_classes
im_height, im_width, _ = image.shape
img = image
scale = [img.shape[1], img.shape[0], img.shape[1], img.shape[0]]
if use_gpu:
images = Variable(preproc(img)[0].unsqueeze(0).cuda(), requires_grad=False)
else:
images = Variable(preproc(img)[0].unsqueeze(0), requires_grad=False)
out = model(images, phase='eval')
# detect
detections = detector.forward(out)
_scores = []
_labels = []
_coords = []
batch = 0
for j in range(1, num_classes):
for det in detections[0][j]:
if det[0] > 0.45:
d = det.cpu().numpy()
score, box = d[0], d[1:]
box *= scale
_labels.append(j - 1)
_coords.append(box)
_scores.append(score)
COLORS = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]
FONT = cv2.FONT_HERSHEY_SIMPLEX
#print("newjson: "+ str(_coords.__len__()))
for label, score, coord in zip(_labels, _scores, _coords):
label_id = load_template_json(label, json_path)
xmin, ymin, xmax, ymax = coord
bbox = {
# {
'x': xmin* 1,
'y': ymin* 1,
'w': (xmax - xmin)* 1,
'h': (ymax - ymin)* 1,
'id': label_id
# },
# 'category': category_index[classes[i]]['name'],
# 'score': float(scores[i])
}
bboxes.append(bbox)
return bboxes
def load_data(cfg, phase):
if phase == 'train':
dataset = dataset_map[cfg.DATASET](cfg.DATASET_DIR, cfg.TRAIN_SETS,
preproc(cfg.IMAGE_SIZE,
cfg.PIXEL_MEANS, cfg.PROB))
data_loader = data.DataLoader(dataset,
cfg.TRAIN_BATCH_SIZE,
num_workers=cfg.NUM_WORKERS,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
if phase == 'eval':
dataset = dataset_map[cfg.DATASET](cfg.DATASET_DIR, cfg.TEST_SETS,
preproc(cfg.IMAGE_SIZE,
cfg.PIXEL_MEANS, -1))
data_loader = data.DataLoader(dataset,
cfg.TEST_BATCH_SIZE,
num_workers=cfg.NUM_WORKERS,
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
if phase == 'test':
dataset = dataset_map[cfg.DATASET](cfg.DATASET_DIR, cfg.TEST_SETS,
preproc(cfg.IMAGE_SIZE,
cfg.PIXEL_MEANS, -2))
data_loader = data.DataLoader(dataset,
cfg.TEST_BATCH_SIZE,
num_workers=cfg.NUM_WORKERS,
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
if phase == 'visualize':
dataset = dataset_map[cfg.DATASET](cfg.DATASET_DIR, cfg.TEST_SETS,
preproc(cfg.IMAGE_SIZE,
cfg.PIXEL_MEANS, 1))
data_loader = data.DataLoader(dataset,
cfg.TEST_BATCH_SIZE,
num_workers=cfg.NUM_WORKERS,
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
if phase == 'correlation':
dataset = dataset_map[cfg.DATASET](cfg.DATASET_DIR, cfg.TRAIN_SETS,
preproc(cfg.IMAGE_SIZE,
cfg.PIXEL_MEANS, -1))
data_loader = data.DataLoader(dataset,
1,
num_workers=cfg.NUM_WORKERS,
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
return data_loader
def _detect_one_image(selfself,
model,
np_image,
preproc,
detector,
use_gpu=True):
model.eval()
#model.onnx_export = True
num_classes = detector.num_classes
img = np_image
scale = [img.shape[1], img.shape[0], img.shape[1], img.shape[0]]
#scale = [img.shape[0], img.shape[1], img.shape[0], img.shape[1]]
if use_gpu:
images = Variable(preproc(img)[0].unsqueeze(0).cuda(),
requires_grad=False)
else:
images = Variable(preproc(img)[0].unsqueeze(0),
requires_grad=False)
out = model(images, phase='eval')
# detect
detections = detector.forward(out)
_scores = []
_labels = []
_coords = []
batch = 0
for j in range(1, num_classes):
for det in detections[0][j]:
if det[0] > 0.45:
d = det.cpu().numpy()
score, box = d[0], d[1:]
box *= scale
_labels.append(j - 1)
_coords.append(box)
_scores.append(score)
COLORS = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]
FONT = cv2.FONT_HERSHEY_SIMPLEX
for label, score, coord in zip(_labels, _scores, _coords):
cv2.rectangle(img, (int(coord[0]), int(coord[1])),
(int(coord[2]), int(coord[3])), COLORS[label % 3], 2)
cv2.putText(
img, '{label}: {score:.3f}'.format(label=label, score=score),
(int(coord[0]), int(coord[1]) + 60), FONT, 0.5,
COLORS[label % 3], 2)
#
#
# cv2.imwrite('/tmp/ddd_result.jpg', image)
print("%d boxes detected" % (len(_labels)))
def visualize_epoch(self, model, data_loader, priorbox, writer, epoch, use_gpu):
model.eval()
img_index = random.randint(0, len(data_loader.dataset)-1)
#img_index = 1
# get img
image = data_loader.dataset.pull_image(img_index)
anno = data_loader.dataset.pull_anno(img_index)
# visualize archor box
viz_prior_box(writer, priorbox, image, epoch)
# get preproc
preproc = data_loader.dataset.preproc
preproc.add_writer(writer, epoch)
# preproc.p = 0.6
# preproc image & visualize preprocess prograss
images = Variable(preproc(image, anno)[0].unsqueeze(0), volatile=True)
if use_gpu:
images = images.cuda()
# visualize feature map in base and extras
base_out = viz_module_feature_maps(writer, model.base, images, module_name='base', epoch=epoch)
extras_out = viz_module_feature_maps(writer, model.extras, base_out, module_name='extras', epoch=epoch)
# visualize feature map in feature_extractors
viz_feature_maps(writer, model(images, 'feature'), module_name='feature_extractors', epoch=epoch)
model.train()
images.requires_grad = True
images.volatile=False
base_out = viz_module_grads(writer, model, model.base, images, images, preproc.means, module_name='base', epoch=epoch)
def __init__(self):
self.cfg = cfg
self.preproc = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS,
-2)
# Build model
print('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = Variable(self.priorbox.forward(), volatile=True)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# if torch.cuda.device_count() > 1:
# self.model = torch.nn.DataParallel(self.model).module
# Print the model architecture and parameters
print('Model architectures:\n{}\n'.format(self.model))
num_parameters = sum([l.nelement() for l in self.model.parameters()])
print('number of parameters: {}'.format(num_parameters))
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# Set the logger
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self, input_cfg=None, root_dir=None, viz_arch=False):
global cfg
self.cfg = cfg if input_cfg is None else input_cfg
cfg = self.cfg
# Build model
logger.info('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = self.priorbox.forward()
# Print the model architecture and parameters
if viz_arch is True:
logger.info('Model architectures:\n{}\n'.format(self.model))
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
#self.half = False
self.half = cfg.MODEL.HALF_PRECISION
if self.half:
self.model = BN_convert_float(self.model.half())
self.priors = self.priors.half()
if self.use_gpu:
logger.info('Utilize GPUs for computation')
logger.info('Number of GPU available: {}'.format(torch.cuda.device_count()))
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# self.model = torch.nn.DataParallel(self.model).module
# Utilize half precision
# Build preprocessor and detector
self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Load weight:
if cfg.RESUME_CHECKPOINT == '':
AssertionError('RESUME_CHECKPOINT can not be empty')
if cfg.RESUME_CHECKPOINT == 'latest':
checkpoint_file = os.path.join(cfg.EXP_DIR, 'checkpoint_list.txt')
assert os.path.exists(checkpoint_file), \
'RESUME_CHECKPOINT set to \'latest\' but {} does not exist'.format(checkpoint_file)
with open(checkpoint_file, 'r') as f:
last_line = f.readlines()[-1]
checkpoint_file = last_line[last_line.find(':') + 2:-1]
else:
checkpoint_file = cfg.RESUME_CHECKPOINT
if root_dir is not None:
checkpoint_file = os.path.join(root_dir, checkpoint_file)
logger.info('=> loading checkpoint {:s}'.format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
self.model.load_state_dict(checkpoint)
# test only
self.model.eval()
def __init__(self, viz_arch=False):
self.cfg = cfg
# Build model
print('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = Variable(self.priorbox.forward(), volatile=True)
# Print the model architecture and parameters
if viz_arch is True:
print('Model architectures:\n{}\n'.format(self.model))
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
self.device = torch.device(
'gpu') if torch.cuda.is_available() else torch.device('cpu')
self.half = False
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# self.model = torch.nn.DataParallel(self.model).module
# Utilize half precision
self.half = cfg.MODEL.HALF_PRECISION
if self.half:
self.model = self.model.half()
self.priors = self.priors.half()
# Build preprocessor and detector
self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE,
cfg.DATASET.PIXEL_MEANS, -2)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Load weight:
if cfg.RESUME_CHECKPOINT == '':
AssertionError('RESUME_CHECKPOINT can not be empty')
print('=> loading checkpoint {:s}'.format(cfg.RESUME_CHECKPOINT))
# checkpoint = torch.load(cfg.RESUME_CHECKPOINT)
checkpoint = torch.load(cfg.RESUME_CHECKPOINT,
map_location='gpu' if self.use_gpu else 'cpu')
self.model.load_state_dict(checkpoint)
# test only
self.model.eval()
def visualize_epoch(self, model, data_loader, priorbox, writer, epoch, use_gpu):
model.eval()
img_index = random.randint(0, len(data_loader.dataset)-1)
# get img
image = data_loader.dataset.pull_image(img_index)
anno = data_loader.dataset.pull_anno(img_index)
# get preproc
preproc = data_loader.dataset.preproc
preproc.add_writer(writer, epoch)
# visualize archor box
viz_prior_box(writer, priorbox, image, epoch)
# preproc image & visualize preprocess prograss
images = Variable(preproc(image, anno)[0].unsqueeze(0), volatile=True)
if use_gpu:
images = images.cuda()
def run_detection(data_dir, coco_gt, im_ids):
model, priorbox = create_model(cfg.MODEL)
priors = Variable(priorbox.forward(), volatile=True)
detector = Detect(cfg.POST_PROCESS, priors)
# Utilize GPUs for computation
model.cuda()
priors.cuda()
cudnn.benchmark = True
preprocess = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2)
resume_checkpoint(model, args.weight)
num_classes = 2
results = []
time_all = []
time_per_step = {
"nms_time": [],
"cpu_tims": [],
"scores_time": [],
"box_time": [],
"gpunms_time": [],
"base_time": [],
"extra_time": [],
"head_time": []
}
for i, index in enumerate(im_ids):
# load img
print('evaluating image {}/{}'.format(i, len(im_ids)))
im_data = coco_gt.loadImgs(ids=index)[0]
img = cv2.imread(os.path.join(data_dir, 'frames',
im_data['file_name']))
scale = [img.shape[1], img.shape[0], img.shape[1], img.shape[0]]
img_shape = img.shape
images = Variable(preprocess(img)[0].unsqueeze(0).cuda(),
volatile=True)
img_dict = {'version': 0, 'time': 0., 'camera_id': 0, 'image': img}
# run detection model
torch.cuda.synchronize()
time_all_start = time.perf_counter()
# forward
out = model(images, phase='eval')
# detect
detections = detector.forward(out)
torch.cuda.synchronize()
time_all_end = time.perf_counter()
time_all.append(1000 * (time_all_end - time_all_start))
scores = []
cls_boxes = []
for det in detections[0][1]:
if det[0] > 0:
d = det.cpu().numpy()
score, box = d[0], d[1:]
box *= scale
scores.append(score)
cls_boxes.append(box)
#print(score)
#print(box)
output_blobs = {}
output_blobs['scores'] = scores
output_blobs['cls_boxes'] = cls_boxes
print(np.array(cls_boxes).shape)
output_dict = postprocess(output_blobs, 1., img_shape)
if len(output_dict['people']) == 0:
continue
# save result
entry_index = 0
for person in output_dict['people']:
entry_result = {
"image_id": index,
"category_id": 1,
"bbox": person['bbox_ltwh'].tolist(),
"score": person['score']
}
results.append(entry_result)
# save results as json file
with open(json_dt, 'w') as f:
json.dump(results, f)
print('detection results saved in {}'.format(json_dt))
print('average running time: {}ms'.format(sum(time_all) / len(time_all)))
