Skip to content

SeetaResearch / SeetaDet

Go to a project
Commit 71593766 by Ting PAN
Options

Add deploy and XML support

1 parent 071996af
Inline Side-by-side
Showing with 274 additions and 67 deletions
lib/datasets/taas.py
......@@ -97,21 +97,54 @@ class TaaS(imdb):
# #
##############################################
def _write_xml_bbox_results(self, all_boxes, gt_recs, output_dir):
from xml.dom import minidom
import xml.etree.ElementTree as ET
ix = 0
for image_id, rec in gt_recs.items():
root = ET.Element('annotation')
ET.SubElement(root, 'filename').text = str(image_id)
for cls_ind, cls in enumerate(self.classes):
if cls == '__background__':
continue
detections = all_boxes[cls_ind][ix]
if len(detections) == 0:
continue
for k in range(detections.shape[0]):
if detections[k, -1] < cfg.VIS_TH:
continue
object = ET.SubElement(root, 'object')
ET.SubElement(object, 'name').text = cls
ET.SubElement(object, 'difficult').text = '0'
bnd_box = ET.SubElement(object, 'bndbox')
ET.SubElement(bnd_box, 'xmin').text = str(detections[k][0])
ET.SubElement(bnd_box, 'ymin').text = str(detections[k][1])
ET.SubElement(bnd_box, 'xmax').text = str(detections[k][2])
ET.SubElement(bnd_box, 'ymax').text = str(detections[k][3])
ix += 1
rawText = ET.tostring(root)
dom = minidom.parseString(rawText)
with open('{}/{}.xml'.format(output_dir, image_id), 'w') as f:
dom.writexml(f, "", "\t", "\n", "utf-8")
def _write_voc_bbox_results(self, all_boxes, gt_recs, output_dir):
for cls_ind, cls in enumerate(self.classes):
if cls == '__background__': continue
if cls == '__background__':
continue
print('Writing {} VOC format bbox results'.format(cls))
filename = self._get_voc_results_T(output_dir).format(cls)
with open(filename, 'wt') as f:
ix = 0
for image_id, rec in gt_recs.items():
dets = all_boxes[cls_ind][ix]; ix += 1
if dets == []: continue
if len(dets) == 0:
continue
for k in range(dets.shape[0]):
f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.
format(image_id, dets[k, -1],
dets[k, 0] + 1, dets[k, 1] + 1,
dets[k, 2] + 1, dets[k, 3] + 1))
f.write(
'{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'
.format(image_id, dets[k, -1],
dets[k, 0] + 1, dets[k, 1] + 1,
dets[k, 2] + 1, dets[k, 3] + 1))
def _write_voc_segm_results(self, all_boxes, all_masks, output_dir):
for cls_inds, cls in enumerate(self.classes):
......@@ -151,40 +184,6 @@ class TaaS(imdb):
##############################################
# #
# ROT #
# #
##############################################
def _write_voc_rbox_results(self, all_boxes, gt_recs, output_dir):
for cls_ind, cls in enumerate(self.classes):
if cls == '__background__': continue
print('Writing {} VOC format rbox results'.format(cls))
filename = self._get_voc_results_T(output_dir).format(cls)
with open(filename, 'wt') as f:
ix = 0
for image_id, rec in gt_recs.items():
dets = all_boxes[cls_ind][ix]; ix += 1
if dets == []: continue
for k in range(dets.shape[0]):
f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f} {:.2f}\n'.
format(image_id, dets[k, -1],
dets[k, 0] + 1, dets[k, 1] + 1,
dets[k, 2] + 1, dets[k, 3] + 1, dets[k, 4]))
def _do_voc_rbox_eval(self, gt_recs, output_dir, IoU=0.5, use_07_metric=True):
aps = []
print('VOC07 metric? ' + ('Yes' if use_07_metric else 'No'))
for i, cls in enumerate(self._classes):
if cls == '__background__': continue
det_file = self._get_voc_results_T(output_dir).format(cls)
rec, prec, ap = voc_rbox_eval(det_file, gt_recs, cls,
IoU=IoU, use_07_metric=use_07_metric)
if ap > 0: aps += [ap]
print('AP for {} = {:.4f}'.format(cls, ap))
print('Mean AP = {:.4f}\n'.format(np.mean(aps)))
##############################################
# #
# COCO #
# #
##############################################
......@@ -325,7 +324,8 @@ class TaaS(imdb):
msks = masks[ix]; ix += 1
keep = filter_boxes(dets)
im_h, im_w = rec['height'], rec['width']
if len(keep) == 0: continue
if len(keep) == 0:
continue
scores = dets[:, -1]
mask_encode = self._encode_coco_masks(
msks[keep], dets[keep, :4], im_h, im_w)
......@@ -425,34 +425,35 @@ class TaaS(imdb):
protocol = cfg.TEST.PROTOCOL
if 'voc' in protocol:
self._write_voc_bbox_results(all_boxes, gt_recs, output_dir)
if not 'wo' in protocol:
print('\n~~~~~~ Evaluation IoU@0.5 ~~~~~~')
self._do_voc_bbox_eval(gt_recs, output_dir,
IoU=0.5, use_07_metric='2007' in protocol)
print('~~~~~~ Evaluation IoU@0.7 ~~~~~~')
self._do_voc_bbox_eval(gt_recs, output_dir,
IoU=0.7, use_07_metric='2007' in protocol)
elif 'rot' in protocol:
self._write_voc_rbox_results(all_boxes, gt_recs, output_dir)
if not 'wo' in protocol:
if 'wo' not in protocol:
print('\n~~~~~~ Evaluation IoU@0.5 ~~~~~~')
self._do_voc_rbox_eval(gt_recs, output_dir,
IoU=0.5, use_07_metric='2007' in protocol)
self._do_voc_bbox_eval(
gt_recs, output_dir, IoU=0.5,
use_07_metric='2007' in protocol)
print('~~~~~~ Evaluation IoU@0.7 ~~~~~~')
self._do_voc_rbox_eval(gt_recs, output_dir,
IoU=0.7, use_07_metric='2007' in protocol)
self._do_voc_bbox_eval(
gt_recs, output_dir, IoU=0.7,
use_07_metric='2007' in protocol)
elif 'xml' in protocol:
if cfg.EXP_DIR != '':
output_dir = cfg.EXP_DIR
self._write_xml_bbox_results(all_boxes, gt_recs, output_dir)
elif 'coco' in protocol:
from lib.pycocotools.coco import COCO
if os.path.exists(cfg.TEST.JSON_FILE):
coco = COCO(cfg.TEST.JSON_FILE)
# We should override category id before writing results
cats = coco.loadCats(coco.getCatIds())
self._class_to_cat_id = dict(
zip([c['name'] for c in cats], coco.getCatIds()))
else: coco = None
res_file = self._write_coco_bbox_results(all_boxes, gt_recs, output_dir)
if not 'wo' in protocol:
if coco is None: coco = COCO(self._write_coco_bbox_annotations(gt_recs, output_dir))
self._class_to_cat_id = dict(zip(
[c['name'] for c in cats], coco.getCatIds()))
else:
coco = None
res_file = self._write_coco_bbox_results(
all_boxes, gt_recs, output_dir)
if 'wo' not in protocol:
if coco is None:
ann_file = self._write_coco_bbox_annotations(gt_recs, output_dir)
coco = COCO(ann_file)
self._do_coco_bbox_eval(coco, res_file)
def evaluate_segmentations(self, all_boxes, all_masks, gt_recs, output_dir):
......@@ -461,11 +462,13 @@ class TaaS(imdb):
self._write_voc_segm_results(all_boxes, all_masks, output_dir)
if not 'wo' in protocol:
print('\n~~~~~~ Evaluation IoU@0.5 ~~~~~~')
self._do_voc_segm_eval(gt_recs, output_dir,
IoU=0.5, use_07_metric='2007' in protocol)
self._do_voc_segm_eval(
gt_recs, output_dir, IoU=0.5,
use_07_metric='2007' in protocol)
print('~~~~~~ Evaluation IoU@0.7 ~~~~~~')
self._do_voc_segm_eval(gt_recs, output_dir,
IoU=0.7, use_07_metric='2007' in protocol)
self._do_voc_segm_eval(
gt_recs, output_dir, IoU=0.7,
use_07_metric='2007' in protocol)
elif 'coco' in protocol:
from lib.pycocotools.coco import COCO
if os.path.exists(cfg.TEST.JSON_FILE):
......@@ -486,4 +489,4 @@ class TaaS(imdb):
self.config['cleanup'] = False
else:
self.config['use_salt'] = True
self.config['cleanup'] = True
\ No newline at end of file
self.config['cleanup'] = True
tools/deploy.py 0 → 100644
# ------------------------------------------------------------
# Copyright (c) 2017-present, SeetaTech, Co.,Ltd.
#
# Licensed under the BSD 2-Clause License.
# You should have received a copy of the BSD 2-Clause License
# along with the software. If not, See,
#
# <https://opensource.org/licenses/BSD-2-Clause>
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import base64
import importlib
import sys
import argparse
import cv2
import numpy as np
import pprint
from seetaas_helper import visualization_test
from lib.core.config import cfg
from lib.core.coordinator import Coordinator
from lib.modeling.detector import Detector
import lib.ssd.test
import lib.faster_rcnn.test
from lib.faster_rcnn.test import nms, soft_nms
##############################################
# #
# ARGS #
# #
##############################################
def parse_args():
"""Parse input arguments"""
parser = argparse.ArgumentParser(description='Test a Detection Network')
parser.add_argument('--cfg', dest='cfg_file',
help='optional config file', default=None, type=str)
parser.add_argument('--exp_dir', dest='exp_dir',
help='experiment dir',
default=None, type=str)
parser.add_argument('--iter', dest='iter', help='global step',
default=0, type=int)
parser.add_argument('--workers', dest='num_workers',
help='number of workers',
default=1, type=int)
if len(sys.argv) == 233:
parser.print_help()
sys.exit(1)
args = parser.parse_args()
return args
##############################################
# #
# UTILS #
# #
##############################################
def get_image(base64_str):
if detector is None:
return {
"state": "False",
"message": "detect model is not init",
"objects": [],
"res": 2,
}
try:
str = base64.b64decode(base64_str)
im = np.fromstring(str, np.uint8)
im = cv2.imdecode(im, cv2.IMREAD_COLOR)
return im
except:
return {
"state": "False",
"message": "detect image is not valid",
"objects": [],
"res": 1,
}
def send_detections(boxes_this_image):
detections = []
for j, name in enumerate(cfg.MODEL.CLASSES):
if name == '__background__':
continue
dets = boxes_this_image[j] # [num, {x1, y1, x2, y2, score}]
keep_inds = np.where(dets[:, 4] > cfg.VIS_TH)[0]
dets = dets[keep_inds]
cls_inds = np.ones((dets.shape[0], 1), dtype=np.float32) * j
dets = np.hstack((dets.astype(np.float32, copy=False), cls_inds))
detections.extend(dets.tolist())
return {"state": "True", "message": "", "objects": detections, "res": 0}
##############################################
# #
# PROCEDURE #
# #
##############################################
def ssd_infer(base64_str):
im = get_image(base64_str)
if not isinstance(im, np.ndarray):
return im
batch_scores, batch_boxes = \
lib.ssd.test.ims_detect(detector, [im])
scores, boxes = batch_scores[0], batch_boxes[0]
boxes_this_image = [[]]
for j in range(1, cfg.MODEL.NUM_CLASSES):
inds = np.where(scores[:, j] > cfg.TEST.SCORE_THRESH)[0]
cls_scores = scores[inds, j]
cls_boxes = boxes[inds]
pre_nms_inds = np.argsort(-cls_scores)[0: cfg.TEST.NMS_TOP_K]
cls_scores = cls_scores[pre_nms_inds]
cls_boxes = cls_boxes[pre_nms_inds]
cls_dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])) \
.astype(np.float32, copy=False)
if cfg.TEST.USE_SOFT_NMS:
keep = soft_nms(
cls_dets, cfg.TEST.NMS,
method=cfg.TEST.SOFT_NMS_METHOD,
sigma=cfg.TEST.SOFT_NMS_SIGMA,
)
else:
keep = nms(
cls_dets, cfg.TEST.NMS,
force_cpu=True,
)
cls_dets = cls_dets[keep, :]
boxes_this_image.append(cls_dets)
return send_detections(boxes_this_image)
def faster_rcnn_infer(base64_str):
im = get_image(base64_str)
if not isinstance(im, np.ndarray):
return im
scores, boxes = lib.faster_rcnn.test.im_detect(detector, im)
boxes_this_image = [[]]
for j in range(1, cfg.MODEL.NUM_CLASSES):
inds = np.where(scores[:, j] > cfg.TEST.SCORE_THRESH)[0]
cls_scores = scores[inds, j]
cls_boxes = boxes[inds, j * 4:(j + 1) * 4]
cls_dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])). \
astype(np.float32, copy=False)
if cfg.TEST.USE_SOFT_NMS:
keep = soft_nms(
cls_dets, cfg.TEST.NMS,
method=cfg.TEST.SOFT_NMS_METHOD,
sigma=cfg.TEST.SOFT_NMS_SIGMA,
)
else:
keep = nms(cls_dets, cfg.TEST.NMS, force_cpu=True)
cls_dets = cls_dets[keep, :]
boxes_this_image.append(cls_dets)
return send_detections(boxes_this_image)
##############################################
# #
# MAIN #
# #
##############################################
@visualization_test.Deploy.register
def infer(base64_str):
infer_procedure = globals()['{}_infer'.format(cfg.MODEL.TYPE)]
return infer_procedure(base64_str)
args = parse_args()
print('Called with args:')
print(args)
coordinator = Coordinator(args.cfg_file, exp_dir=args.exp_dir)
print('Using config:')
pprint.pprint(cfg)
checkpoint = coordinator.checkpoint(global_step=args.iter, wait=False)
test_engine = importlib.import_module('lib.{}.test'.format(cfg.MODEL.TYPE))
detector = Detector().eval().cuda(cfg.GPU_ID)
detector.load_weights(checkpoint)
detector.optimize_for_inference()
# setup database
visualization_test.Deploy.run(args.num_workers)
\ No newline at end of file
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!