Skip to content

SeetaResearch / SeetaDet

Go to a project
Commit f8359d17 by Ting PAN
Options

Adapt to SeetaRecord

1 parent ca255ea0
Inline Side-by-side
Showing with 464 additions and 1232 deletions
CHANGES
------------------------------------------------------------------------
The list of most significant changes made over time in SeetaDet.
SeetaDet 0.2.0 (20190929)
Dragon Minimum Required (Version 0.3.0.dev20190929)
Changes:
Preview Features:
- Use SeetaRecord instead of LMDB.
- Flatten the implementation of layers.
Bugs fixed:
- None
------------------------------------------------------------------------
SeetaDet 0.1.2 (20190723)
Dragon Minimum Required (Version 0.3.0.0)
......
compile/make.sh
#!/bin/sh
# delete cache
rm -r build install *.c *.cpp
# compile proto files
protoc -I ../lib/proto --python_out=../lib/proto ../lib/proto/anno.proto
# compile cython modules
python setup.py build_ext --inplace
# compile cuda modules
cd build
cmake .. && make install && cd ..
cd build && cmake .. && make install && cd ..
# setup
cp -r install/lib ../
configs/faster_rcnn/coco_faster_rcnn_R-101-FPN_1x.yml
......@@ -32,15 +32,15 @@ FRCNN:
ROI_XFORM_METHOD: RoIAlign
ROI_XFORM_RESOLUTION: 7
TRAIN:
WEIGHTS: '/data/models/imagenet/R-101.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k_lmdb'
WEIGHTS: '/model/R-101.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k'
IMS_PER_BATCH: 2
USE_DIFF: False # Do not use crowd objects
BATCH_SIZE: 512
SCALES: [800]
MAX_SIZE: 1333
TEST:
DATABASE: '/data/coco_2014_minival_lmdb'
DATABASE: '/data/coco_2014_minival'
JSON_FILE: '/data/instances_minival2014.json'
PROTOCOL: 'coco'
RPN_POST_NMS_TOP_N: 1000
......
configs/faster_rcnn/coco_faster_rcnn_R-101-FPN_2x.yml
......@@ -32,15 +32,15 @@ FRCNN:
ROI_XFORM_METHOD: RoIAlign
ROI_XFORM_RESOLUTION: 7
TRAIN:
WEIGHTS: '/data/models/imagenet/R-101.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k_lmdb'
WEIGHTS: '/model/R-101.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k'
IMS_PER_BATCH: 2
USE_DIFF: False # Do not use crowd objects
BATCH_SIZE: 512
SCALES: [800]
MAX_SIZE: 1333
TEST:
DATABASE: '/data/coco_2014_minival_lmdb'
DATABASE: '/data/coco_2014_minival'
JSON_FILE: '/data/instances_minival2014.json'
PROTOCOL: 'coco'
RPN_POST_NMS_TOP_N: 1000
......
configs/faster_rcnn/voc_faster_rcnn_R-50-FPN.yml
......@@ -23,14 +23,14 @@ FRCNN:
ROI_XFORM_METHOD: RoIAlign
ROI_XFORM_RESOLUTION: 7
TRAIN:
WEIGHTS: '/data/models/imagenet/R-50.Affine.pth'
DATABASE: '/data/voc_0712_trainval_lmdb'
WEIGHTS: '/model/R-50.Affine.pth'
DATABASE: '/data/voc_0712_trainval'
IMS_PER_BATCH: 2
BATCH_SIZE: 128
SCALES: [600]
MAX_SIZE: 1000
TEST:
DATABASE: '/data/voc_2007_test_lmdb'
DATABASE: '/data/voc_2007_test'
PROTOCOL: 'voc2007' # 'voc2007', 'voc2010', 'coco'
RPN_POST_NMS_TOP_N: 1000
SCALES: [600]
......
configs/faster_rcnn/voc_faster_rcnn_VGG-16-C4.yml
......@@ -28,15 +28,15 @@ FRCNN:
ROI_XFORM_RESOLUTION: 7
MLP_HEAD_DIM: 4096
TRAIN:
WEIGHTS: '/data/models/imagenet/VGG16.RCNN.pth'
DATABASE: '/data/voc_0712_trainval_lmdb'
WEIGHTS: '/model/VGG16.RCNN.pth'
DATABASE: '/data/voc_0712_trainval'
RPN_MIN_SIZE: 16
IMS_PER_BATCH: 2
BATCH_SIZE: 128
SCALES: [600]
MAX_SIZE: 1000
TEST:
DATABASE: '/data/voc_2007_test_lmdb'
DATABASE: '/data/voc_2007_test'
PROTOCOL: 'voc2007' # 'voc2007', 'voc2010', 'coco'
RPN_MIN_SIZE: 16
RPN_POST_NMS_TOP_N: 300
......
configs/retinanet/coco_retinanet_400_R-50-FPN_1x.yml
......@@ -32,13 +32,13 @@ FPN:
RPN_MIN_LEVEL: 3
RPN_MAX_LEVEL: 7
TRAIN:
WEIGHTS: '/data/models/imagenet/R-50.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k_lmdb'
WEIGHTS: '/model/R-50.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k'
IMS_PER_BATCH: 8
SCALES: [400]
MAX_SIZE: 666
TEST:
DATABASE: '/data/coco_2014_minival_lmdb'
DATABASE: '/data/coco_2014_minival'
JSON_FILE: '/data/instances_minival2014.json'
PROTOCOL: 'coco'
IMS_PER_BATCH: 1
......
configs/retinanet/coco_retinanet_400_R-50-FPN_4x.yml
......@@ -36,8 +36,8 @@ DROPBLOCK:
DROP_ON: True
DECREMENT: 0.000005 # * 20000 = 0.1
TRAIN:
WEIGHTS: '/data/models/imagenet/R-50.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k_lmdb'
WEIGHTS: '/model/R-50.Affine.pth'
DATABASE: '/data/coco_2014_trainval35k'
IMS_PER_BATCH: 8
SCALES: [400]
MAX_SIZE: 666
......@@ -45,7 +45,7 @@ TRAIN:
COLOR_JITTERING: True
SCALE_RANGE: [0.75, 1.33]
TEST:
DATABASE: '/data/coco_2014_minival_lmdb'
DATABASE: '/data/coco_2014_minival'
JSON_FILE: '/data/instances_minival2014.json'
PROTOCOL: 'coco'
IMS_PER_BATCH: 1
......
configs/retinanet/voc_retinanet_300_AirNet-FPN.yml
......@@ -23,8 +23,8 @@ FPN:
RPN_MIN_LEVEL: 3
RPN_MAX_LEVEL: 7
TRAIN:
WEIGHTS: '/data/models/imagenet/AirNet.Affine.pth'
DATABASE: '/data/voc_0712_trainval_lmdb'
WEIGHTS: '/model/AirNet.Affine.pth'
DATABASE: '/data/voc_0712_trainval'
IMS_PER_BATCH: 32
SCALES: [300]
MAX_SIZE: 500
......@@ -32,7 +32,7 @@ TRAIN:
SCALE_JITTERING: True
COLOR_JITTERING: True
TEST:
DATABASE: '/data/voc_2007_test_lmdb'
DATABASE: '/data/voc_2007_test'
PROTOCOL: 'voc2007' # 'voc2007', 'voc2010', 'coco'
IMS_PER_BATCH: 1
SCALES: [300]
......
configs/retinanet/voc_retinanet_300_R-18-FPN.yml
......@@ -24,8 +24,8 @@ FPN:
RPN_MIN_LEVEL: 3
RPN_MAX_LEVEL: 7
TRAIN:
WEIGHTS: '/data/models/imagenet/R-18.Affine.pth'
DATABASE: '/data/voc_0712_trainval_lmdb'
WEIGHTS: '/model/R-18.Affine.pth'
DATABASE: '/data/voc_0712_trainval'
IMS_PER_BATCH: 32
SCALES: [300]
MAX_SIZE: 500
......@@ -33,7 +33,7 @@ TRAIN:
SCALE_JITTERING: True
COLOR_JITTERING: True
TEST:
DATABASE: '/data/voc_2007_test_lmdb'
DATABASE: '/data/voc_2007_test'
PROTOCOL: 'voc2007' # 'voc2007', 'voc2010', 'coco'
IMS_PER_BATCH: 1
SCALES: [300]
......
configs/retinanet/voc_retinanet_300_R-34-FPN.yml
......@@ -24,8 +24,8 @@ FPN:
RPN_MIN_LEVEL: 3
RPN_MAX_LEVEL: 7
TRAIN:
WEIGHTS: '/data/models/imagenet/R-34.Affine.pth'
DATABASE: '/data/voc_0712_trainval_lmdb'
WEIGHTS: '/model/R-34.Affine.pth'
DATABASE: '/data/voc_0712_trainval'
IMS_PER_BATCH: 32
SCALES: [300]
MAX_SIZE: 500
......@@ -33,7 +33,7 @@ TRAIN:
SCALE_JITTERING: True
COLOR_JITTERING: True
TEST:
DATABASE: '/data/voc_2007_test_lmdb'
DATABASE: '/data/voc_2007_test'
PROTOCOL: 'voc2007' # 'voc2007', 'voc2010', 'coco'
IMS_PER_BATCH: 1
SCALES: [300]
......
configs/ssd/voc_ssd_300_AirNet-5b.yml
......@@ -29,11 +29,11 @@ SSD:
STRIDES: [8, 16, 32]
ASPECT_RATIOS: [[1, 2, 0.5], [1, 2, 0.5], [1, 2, 0.5]]
TRAIN:
WEIGHTS: '/data/models/imagenet/AirNet.Affine.pth'
DATABASE: '/data/voc_0712_trainval_lmdb'
WEIGHTS: '/model/AirNet.Affine.pth'
DATABASE: '/data/voc_0712_trainval'
IMS_PER_BATCH: 32
TEST:
DATABASE: '/data/voc_2007_test_lmdb'
DATABASE: '/data/voc_2007_test'
PROTOCOL: 'voc2007' # 'voc2007', 'voc2010', 'coco'
IMS_PER_BATCH: 8
NMS_TOP_K: 400
......
configs/ssd/voc_ssd_300_VGG-16.yml
......@@ -32,11 +32,11 @@ SSD:
ASPECT_RATIOS: [[1, 2, 0.5], [1, 2, 0.5, 3, 0.33], [1, 2, 0.5, 3, 0.33],
[1, 2, 0.5, 3, 0.33], [1, 2, 0.5], [1, 2, 0.5]]
TRAIN:
WEIGHTS: '/data/models/imagenet/VGG16.SSD.pth'
DATABASE: '/data/voc_0712_trainval_lmdb'
WEIGHTS: '/model/VGG16.SSD.pth'
DATABASE: '/data/voc_0712_trainval'
IMS_PER_BATCH: 32
TEST:
DATABASE: '/data/voc_2007_test_lmdb'
DATABASE: '/data/voc_2007_test'
PROTOCOL: 'voc2007' # 'voc2007', 'voc2010', 'coco'
IMS_PER_BATCH: 8
NMS_TOP_K: 400
......
database/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
database/frcnn/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
database/frcnn/pascal_voc/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
database/frcnn/pascal_voc/make_lmdb.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 sys
import os.path as osp
sys.path.insert(0, '../../../')
from database.frcnn.utils.make_from_xml import make_db
if __name__ == '__main__':
VOC_ROOT_DIR = '/home/workspace/datasets/VOC'
# train database: voc_2007_trainval + voc_2012_trainval
make_db(database_file=osp.join(VOC_ROOT_DIR, 'cache/voc_0712_trainval_lmdb'),
images_path=[osp.join(VOC_ROOT_DIR, 'VOCdevkit2007/VOC2007/JPEGImages'),
osp.join(VOC_ROOT_DIR, 'VOCdevkit2012/VOC2012/JPEGImages')],
annotations_path=[osp.join(VOC_ROOT_DIR, 'VOCdevkit2007/VOC2007/Annotations'),
osp.join(VOC_ROOT_DIR, 'VOCdevkit2012/VOC2012/Annotations')],
imagesets_path=[osp.join(VOC_ROOT_DIR, 'VOCdevkit2007/VOC2007/ImageSets/Main'),
osp.join(VOC_ROOT_DIR, 'VOCdevkit2012/VOC2012/ImageSets/Main')],
splits=['trainval', 'trainval'])
# test database: voc_2007_test
make_db(database_file=osp.join(VOC_ROOT_DIR, 'cache/voc_2007_test_lmdb'),
images_path=osp.join(VOC_ROOT_DIR, 'VOCdevkit2007/VOC2007/JPEGImages'),
annotations_path=osp.join(VOC_ROOT_DIR, 'VOCdevkit2007/VOC2007/Annotations'),
imagesets_path=osp.join(VOC_ROOT_DIR, 'VOCdevkit2007/VOC2007/ImageSets/Main'),
splits=['test'])
database/frcnn/utils/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
database/frcnn/utils/make_from_dict.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 os
import sys
import time
import cv2
from dragon.tools.db import LMDB
sys.path.insert(0, '../../..')
from lib.proto import anno_pb2 as pb
ZFILL = 8
ENCODE_QUALITY = 95
def set_zfill(value):
global ZFILL
ZFILL = value
def set_quality(value):
global ENCODE_QUALITY
ENCODE_QUALITY = value
def make_datum(image_id, image_file, objects):
anno_datum = pb.AnnotatedDatum()
datum = pb.Datum()
im = cv2.imread(image_file)
datum.height, datum.width, datum.channels = im.shape
datum.encoded = ENCODE_QUALITY != 100
if datum.encoded:
result, im = cv2.imencode('.jpg', im, [int(cv2.IMWRITE_JPEG_QUALITY), ENCODE_QUALITY])
datum.data = im.tostring()
anno_datum.datum.CopyFrom(datum)
anno_datum.filename = image_id
for ix, obj in enumerate(objects):
anno = pb.Annotation()
anno.x1, anno.y1, anno.x2, anno.y2 = obj['bbox']
anno.name = obj['name']
anno.difficult = obj['difficult']
anno_datum.annotation.add().CopyFrom(anno)
return anno_datum
def make_db(database_file, images_path, gt_recs, ext='.png'):
if os.path.isdir(database_file) is True:
raise ValueError('The database path is already exist.')
else:
root_dir = database_file[:database_file.rfind('/')]
if not os.path.exists(root_dir):
os.makedirs(root_dir)
print('Start Time: ', time.strftime("%a, %d %b %Y %H:%M:%S", time.gmtime()))
db = LMDB(max_commit=10000)
db.open(database_file, mode='w')
count = 0
total_line = len(gt_recs)
start_time = time.time()
zfill_flag = '{0:0%d}' % (ZFILL)
for image_id, objects in gt_recs.items():
count += 1
if count % 10000 == 0:
now_time = time.time()
print('{0} / {1} in {2:.2f} sec'.format(
count, total_line, now_time - start_time))
db.commit()
image_file = os.path.join(images_path, image_id + ext)
datum = make_datum(image_id, image_file, objects)
db.put(zfill_flag.format(count - 1), datum.SerializeToString())
now_time = time.time()
print('{0} / {1} in {2:.2f} sec'.format(count, total_line, now_time - start_time))
db.commit()
db.close()
end_time = time.time()
print('{0} images have been stored in the database.'.format(total_line))
print('This task finishes within {0:.2f} seconds.'.format(end_time - start_time))
print('The size of database is {0} MB.'.format(
float(os.path.getsize(database_file + '/data.mdb') / 1000 / 1000)))
\ No newline at end of file
database/frcnn/utils/make_from_xml.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 os
import sys
import time
import cv2
import xml.etree.ElementTree as ET
from dragon.tools.db import LMDB
sys.path.insert(0, '../../..')
from lib.proto import anno_pb2 as pb
ZFILL = 8
ENCODE_QUALITY = 95
def set_zfill(value):
global ZFILL
ZFILL = value
def set_quality(value):
global ENCODE_QUALITY
ENCODE_QUALITY = value
def make_datum(image_file, xml_file):
tree = ET.parse(xml_file)
filename = os.path.split(xml_file)[-1]
objs = tree.findall('object')
anno_datum = pb.AnnotatedDatum()
datum = pb.Datum()
im = cv2.imread(image_file)
if im is None or im.shape[0] == 0 or im.shape[1] == 0:
print("XML have not objects ignored: ", xml_file)
return None
datum.height, datum.width, datum.channels = im.shape
datum.encoded = ENCODE_QUALITY != 100
if datum.encoded:
result, im = cv2.imencode('.jpg', im, [int(cv2.IMWRITE_JPEG_QUALITY), ENCODE_QUALITY])
if im is None or im.shape[0] == 0 or im.shape[1] == 0:
print("XML have not objects ignored: ", xml_file)
return None
datum.data = im.tostring()
anno_datum.datum.CopyFrom(datum)
anno_datum.filename = filename.split('.')[0]
if len(objs) == 0:
return None
for ix, obj in enumerate(objs):
anno = pb.Annotation()
bbox = obj.find('bndbox')
x1 = float(bbox.find('xmin').text)
y1 = float(bbox.find('ymin').text)
x2 = float(bbox.find('xmax').text)
y2 = float(bbox.find('ymax').text)
cls = obj.find('name').text.strip()
anno.x1, anno.y1, anno.x2, anno.y2 = (x1, y1, x2, y2)
anno.name = cls
class_name_set.add(cls)
anno.difficult = False
if obj.find('difficult') is not None:
anno.difficult = int(obj.find('difficult').text) == 1
anno_datum.annotation.add().CopyFrom(anno)
return anno_datum
def make_db(
database_file,
images_path,
annotations_path,
imagesets_path,
splits,
):
if os.path.isdir(database_file) is True:
print('Warning: The database path is already exist.')
else:
root_dir = database_file[:database_file.rfind('/')]
if not os.path.exists(root_dir):
os.makedirs(root_dir)
if not isinstance(images_path, list):
images_path = [images_path]
if not isinstance(annotations_path, list):
annotations_path = [annotations_path]
if not isinstance(imagesets_path, list):
imagesets_path = [imagesets_path]
assert len(splits) == len(imagesets_path)
assert len(splits) == len(images_path)
assert len(splits) == len(annotations_path)
print('Start Time: ', time.strftime("%a, %d %b %Y %H:%M:%S", time.gmtime()))
db = LMDB(max_commit=1000)
db.open(database_file, mode='w')
count = 0
total_line = 0
start_time = time.time()
zfill_flag = '{0:0%d}' % ZFILL
for db_idx, split in enumerate(splits):
split_file = os.path.join(imagesets_path[db_idx], split + '.txt')
assert os.path.exists(split_file)
with open(split_file, 'r') as f:
lines = f.readlines()
total_line += len(lines)
for line in lines:
filename = line.strip()
image_file = os.path.join(images_path[db_idx], filename + '.jpg')
xml_file = os.path.join(annotations_path[db_idx], filename + '.xml')
datum = make_datum(image_file, xml_file)
if datum is not None:
count += 1
db.put(zfill_flag.format(count - 1), datum.SerializeToString())
if count % 1000 == 0:
now_time = time.time()
print('{0} / {1} in {2:.2f} sec'.format(
count, total_line, now_time - start_time))
db.commit()
now_time = time.time()
print('{0} / {1} in {2:.2f} sec'.format(count, total_line, now_time - start_time))
db.commit()
db.close()
end_time = time.time()
print('{0} images have been stored in the database.'.format(total_line))
print('This task finishes within {0:.2f} seconds.'.format(end_time - start_time))
print('The size of database is {0} MB.'.format(
float(os.path.getsize(database_file + '/data.mdb') / 1000 / 1000)))
lib/core/config.py
......@@ -68,6 +68,8 @@ __C.TRAIN.BG_THRESH_LO = 0.0
# Use shuffle after each epoch
__C.TRAIN.USE_SHUFFLE = True
# The number of chunks to shuffle
__C.TRAIN.NUM_SHUFFLE_CHUNKS = 0
# Use horizontally-flipped images during training?
__C.TRAIN.USE_FLIPPED = True
......
lib/core/test.py
......@@ -13,16 +13,16 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import multiprocessing as mp
import os
import cv2
from multiprocessing import Queue
from collections import OrderedDict
import dragon
from lib.core.config import cfg
from lib.datasets.factory import get_imdb
# All detectors share the same reader/transformer during testing
from lib.faster_rcnn.data.data_reader import DataReader
from lib.faster_rcnn.data.data_transformer import DataTransformer
from lib.faster_rcnn.data_transformer import DataTransformer
class TestServer(object):
......@@ -31,11 +31,12 @@ class TestServer(object):
self.imdb.competition_mode(cfg.TEST.COMPETITION_MODE)
self.num_images, self.num_classes, self.classes = \
self.imdb.num_images, self.imdb.num_classes, self.imdb.classes
self.data_reader = DataReader(**{'source': self.imdb.source})
self.data_reader = dragon.io.DataReader(
dataset=lambda: dragon.io.SeetaRecordDataset(self.imdb.source))
self.data_transformer = DataTransformer()
self.data_reader.q_out = Queue(cfg.TEST.IMS_PER_BATCH)
self.data_reader.q_out = mp.Queue(cfg.TEST.IMS_PER_BATCH)
self.data_reader.start()
self.gt_recs = OrderedDict()
self.gt_recs = collections.OrderedDict()
self.output_dir = output_dir
if cfg.VIS_ON_FILE:
self.vis_dir = os.path.join(self.output_dir, 'vis')
......@@ -46,9 +47,9 @@ class TestServer(object):
self.data_transformer = transformer_cls()
def get_image(self):
serialized = self.data_reader.q_out.get()
image = self.data_transformer.get_image(serialized)
image_id, objects = self.data_transformer.get_annotations(serialized)
example = self.data_reader.q_out.get()
image = self.data_transformer.get_image(example)
image_id, objects = self.data_transformer.get_annotations(example)
self.gt_recs[image_id] = {
'objects': objects,
'width': image.shape[1],
......@@ -70,11 +71,18 @@ class TestServer(object):
def evaluate_detections(self, all_boxes):
self.imdb.evaluate_detections(
all_boxes, self.get_records(), self.output_dir)
all_boxes,
self.get_records(),
self.output_dir,
)
def evaluate_segmentations(self, all_boxes, all_masks):
self.imdb.evaluate_segmentations(
all_boxes, all_masks, self.get_records(), self.output_dir)
all_boxes,
all_masks,
self.get_records(),
self.output_dir,
)
class InferServer(object):
......@@ -85,7 +93,7 @@ class InferServer(object):
self.num_images, self.num_classes, self.classes = \
len(self.images), cfg.MODEL.NUM_CLASSES, cfg.MODEL.CLASSES
self.data_transformer = DataTransformer()
self.gt_recs = OrderedDict()
self.gt_recs = collections.OrderedDict()
self.output_dir = output_dir
self.image_idx = 0
if cfg.VIS_ON_FILE:
......@@ -101,10 +109,7 @@ class InferServer(object):
image_id = image_name.split('.')[0]
image = cv2.imread(os.path.join(self.images_dir, image_name))
self.image_idx = (self.image_idx + 1) % self.num_images
self.gt_recs[image_id] = {
'width': image.shape[1],
'height': image.shape[0],
}
self.gt_recs[image_id] = {'width': image.shape[1], 'height': image.shape[0]}
return image_id, image
def get_save_filename(self, image_id, ext='.jpg'):
......
lib/datasets/imdb.py
......@@ -14,7 +14,7 @@
# ------------------------------------------------------------
import os
from dragon.tools.db import LMDB
import dragon
from lib.core.config import cfg
......@@ -46,19 +46,18 @@ class imdb(object):
@property
def source(self):
excepted_source = os.path.join(self.cache_path, self.name + '_lmdb')
excepted_source = os.path.join(self.cache_path, self.name)
if not os.path.exists(excepted_source):
raise RuntimeError('Excepted LMDB source from: {}, '
'but it is not existed.'.format(excepted_source))
raise RuntimeError(
'Excepted source from: {}, '
'but it is not existed.'
.format(excepted_source)
)
return excepted_source
@property
def num_images(self):
self._db = LMDB()
self._db.open(self.source)
num_entries = self._db.num_entries()
self._db.close()
return num_entries
return dragon.io.SeetaRecordDataset(self.source).size
def evaluate_detections(self, all_boxes, gt_recs, output_dir):
pass
......
lib/datasets/taas.py
......@@ -17,22 +17,24 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import json
import os
import sys
import json
import numpy as np
import uuid
import cv2
import numpy as np
try:
import cPickle
except:
import pickle as cPickle
from .imdb import imdb
from .voc_eval import voc_bbox_eval, voc_segm_eval
from lib.core.config import cfg
from lib.utils import boxes as box_utils
from lib.datasets.imdb import imdb
from lib.datasets.voc_eval import voc_bbox_eval
from lib.datasets.voc_eval import voc_segm_eval
from lib.pycocotools.mask import encode as encode_masks
from lib.utils import boxes as box_utils
class TaaS(imdb):
......@@ -49,8 +51,11 @@ class TaaS(imdb):
def source(self):
excepted_source = self._source
if not os.path.exists(excepted_source):
raise RuntimeError('Excepted LMDB source from: {}, '
'but it is not existed.'.format(excepted_source))
raise RuntimeError(
'Excepted source from: {}, '
'but it is not existed.'
.format(excepted_source)
)
return excepted_source
##############################################
......
lib/datasets/voc_eval.py
......@@ -28,7 +28,7 @@ except:
from lib.core.config import cfg
from lib.pycocotools.mask_utils import mask_rle2im
from lib.utils.boxes import expand_boxes
from lib.utils.mask_transform import mask_overlap
from lib.utils.mask import mask_overlap
def voc_ap(rec, prec, use_07_metric=False):
......
lib/faster_rcnn/__init__.py
......@@ -13,7 +13,7 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from lib.faster_rcnn.layers.anchor_target_layer import AnchorTargetLayer
from lib.faster_rcnn.layers.data_layer import DataLayer
from lib.faster_rcnn.layers.proposal_layer import ProposalLayer
from lib.faster_rcnn.layers.proposal_target_layer import ProposalTargetLayer
from lib.faster_rcnn.anchor_target_layer import AnchorTargetLayer
from lib.faster_rcnn.data_layer import DataLayer
from lib.faster_rcnn.proposal_layer import ProposalLayer
from lib.faster_rcnn.proposal_target_layer import ProposalTargetLayer
lib/faster_rcnn/layers/anchor_target_layer.py → lib/faster_rcnn/anchor_target_layer.py
......@@ -166,10 +166,10 @@ class AnchorTargetLayer(torch.nn.Module):
bbox_targets = np.zeros((num_inside, 4), dtype=np.float32)
bbox_targets[fg_inds, :] = bbox_transform(
ex_rois=anchors[fg_inds, :],
gt_rois=gt_boxes[argmax_overlaps[fg_inds], 0:4],
gt_rois=gt_boxes[argmax_overlaps[fg_inds], :4],
)
bbox_inside_weights = np.zeros((num_inside, 4), dtype=np.float32)
bbox_inside_weights[labels == 1, :] = np.array((1.0, 1.0, 1.0, 1.0))
bbox_inside_weights[labels == 1, :] = np.array((1., 1., 1., 1.))
bbox_outside_weights = np.zeros((num_inside, 4), dtype=np.float32)
bbox_outside_weights[labels == 1, :] = np.ones((1, 4)) / cfg.TRAIN.RPN_BATCHSIZE
bbox_outside_weights[labels == 0, :] = np.ones((1, 4)) / cfg.TRAIN.RPN_BATCHSIZE
......
lib/faster_rcnn/data/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
lib/faster_rcnn/data/blob_fetcher.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 multiprocessing
import numpy as np
from lib.core.config import cfg
from lib.utils.blob import im_list_to_blob
class BlobFetcher(multiprocessing.Process):
def __init__(self, **kwargs):
super(BlobFetcher, self).__init__()
self.q1_in = self.q2_in = self.q_out = None
self.daemon = True
def get(self, Q_in):
processed_ims, ims_info, all_boxes = [], [], []
for ix in range(cfg.TRAIN.IMS_PER_BATCH):
im, im_scale, gt_boxes = Q_in.get()
processed_ims.append(im)
ims_info.append(list(im.shape[0:2]) + [im_scale])
# Encode boxes by adding the idx of images
im_boxes = np.zeros((gt_boxes.shape[0], gt_boxes.shape[1] + 1), dtype=np.float32)
im_boxes[:, 0:gt_boxes.shape[1]] = gt_boxes
im_boxes[:, -1] = ix
all_boxes.append(im_boxes)
return {
'data': im_list_to_blob(processed_ims),
'ims_info': np.array(ims_info, dtype=np.float32),
'gt_boxes': np.concatenate(all_boxes, axis=0),
}
def run(self):
while True:
if self.q1_in.qsize() >= cfg.TRAIN.IMS_PER_BATCH:
self.q_out.put(self.get(self.q1_in))
elif self.q2_in.qsize() >= cfg.TRAIN.IMS_PER_BATCH:
self.q_out.put(self.get(self.q2_in))
lib/faster_rcnn/data/data_reader.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 math
import multiprocessing
import numpy
from dragon.tools import db
from lib.core.config import cfg
class DataReader(multiprocessing.Process):
"""Collect encoded str from `LMDB`_.
Partition and shuffle records over distributed nodes.
Parameters
----------
source : str
The path of database.
shuffle : bool, optional, default=False
Whether to shuffle the data.
num_chunks : int, optional, default=2048
The number of chunks to split.
"""
def __init__(self, **kwargs):
"""Create a DataReader."""
super(DataReader, self).__init__()
self._source = kwargs.get('source', '')
self._use_shuffle = kwargs.get('shuffle', False)
self._num_chunks = kwargs.get('num_chunks', 2048)
self._part_idx, self._num_parts = 0, 1
self._cursor, self._chunk_cursor = 0, 0
self._chunk_size, self._perm_size = 0, 0
self._head, self._tail, self._num_entries = 0, 0, 0
self._db, self._zfill, self._perm = None, None, None
self._rng_seed = cfg.RNG_SEED
self.q_out = None
self.daemon = True
def element(self):
"""Get the value of current record.
Returns
-------
str
The encoded str.
"""
return self._db.value()
def redirect(self, target):
"""Redirect to the target position.
Parameters
----------
target : int
The key of the record.
Notes
-----
The redirection reopens the database.
You can drop caches by ``echo 3 > /proc/sys/vm/drop_caches``.
This will disturb getting stuck when *Database Size* >> *RAM Size*.
"""
self._db.close()
self._db.open(self._source)
self._cursor = target
self._db.set(str(target).zfill(self._zfill))
def reset(self):
"""Reset the cursor and environment."""
if self._num_parts > 1 or self._use_shuffle:
self._chunk_cursor = 0
self._part_idx = (self._part_idx + 1) % self._num_parts
if self._use_shuffle:
self._perm = numpy.random.permutation(self._perm_size)
self._head = self._part_idx * self._perm_size + self._perm[self._chunk_cursor]
self._tail = self._head * self._chunk_size
if self._head >= self._num_entries: self.next_chunk()
self._tail = self._head + self._chunk_size
self._tail = min(self._num_entries, self._tail)
else:
self._head, self._tail = 0, self._num_entries
self.redirect(self._head)
def next_record(self):
"""Step the cursor of records."""
self._db.next()
self._cursor += 1
def next_chunk(self):
"""Step the cursor of chunks."""
self._chunk_cursor += 1
if self._chunk_cursor >= self._perm_size:
self.reset()
else:
self._head = self._part_idx * self._perm_size + self._perm[self._chunk_cursor]
self._head = self._head * self._chunk_size
if self._head >= self._num_entries:
self.next_chunk()
else:
self._tail = self._head + self._chunk_size
self._tail = min(self._num_entries, self._tail)
self.redirect(self._head)
def run(self):
"""Start the process."""
# Fix seed
numpy.random.seed(self._rng_seed)
# Init db
self._db = db.LMDB()
self._db.open(self._source)
self._zfill = self._db.zfill()
self._num_entries = self._db.num_entries()
epoch_size = self._num_entries // self._num_parts + 1
if self._use_shuffle:
if self._num_chunks <= 0:
# Each chunk has at most 1 record (Record-Wise)
self._chunk_size, self._perm_size = 1, epoch_size
else:
# Search a optimal chunk size (Chunk-Wise)
min_size, max_size = \
1, self._db._total_size * 1.0 \
/ (self._num_chunks * (1 << 20))
while min_size * 2 < max_size: min_size *= 2
self._perm_size = int(math.ceil(
self._db._total_size * 1.1 /
(self._num_parts * min_size << 20)))
self._chunk_size = int(
self._num_entries * 1.0 /
(self._perm_size * self._num_parts) + 1)
limit = (self._num_parts - 0.5) * self._perm_size * self._chunk_size
if self._num_entries <= limit:
# Roll back to Record-Wise shuffle
self._chunk_size, self._perm_size = 1, epoch_size
else:
# One chunk has at most K records
self._chunk_size, self._perm_size = epoch_size, 1
self._perm = numpy.arange(self._perm_size)
# Init env
self.reset()
# Run!
while True:
self.q_out.put(self.element())
self.next_record()
if self._cursor >= self._tail:
if self._num_parts > 1 or self._use_shuffle:
self.next_chunk()
else:
self.reset()
lib/faster_rcnn/data/data_batch.py → lib/faster_rcnn/data_layer.py
......@@ -13,55 +13,70 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from multiprocessing import Queue
import multiprocessing as mp
import time
import dragon
import pprint
import dragon.vm.torch as torch
import numpy as np
from lib.core.config import cfg
from lib.faster_rcnn.data.data_reader import DataReader
from lib.faster_rcnn.data.data_transformer import DataTransformer
from lib.faster_rcnn.data.blob_fetcher import BlobFetcher
from lib.faster_rcnn.data_transformer import DataTransformer
from lib.datasets.factory import get_imdb
from lib.utils import logger
from lib.utils.blob import im_list_to_blob
class DataBatch(object):
"""DataBatch aims to prefetch data by ``Triple-Buffering``.
class DataLayer(torch.nn.Module):
"""Generate a mini-batch of data."""
It takes full advantages of the Process/Thread of Python,
def __init__(self):
super(DataLayer, self).__init__()
database = get_imdb(cfg.TRAIN.DATABASE)
self.data_batch = DataBatch(**{
'dataset': lambda: dragon.io.SeetaRecordDataset(database.source),
'classes': database.classes,
'shuffle': cfg.TRAIN.USE_SHUFFLE,
'num_chunks': cfg.TRAIN.NUM_SHUFFLE_CHUNKS,
'batch_size': cfg.TRAIN.IMS_PER_BATCH * 2,
})
which provides remarkable I/O speed up for scalable distributed training.
def forward(self):
# Get an array blob from the Queue
outputs = self.data_batch.get()
# Zero-Copy the array to tensor
outputs['data'] = torch.from_numpy(outputs['data'])
return outputs
class DataBatch(mp.Process):
"""Prefetch the batch of data."""
"""
def __init__(self, **kwargs):
"""Construct a ``DataBatch``.
Parameters
----------
source : str
The path of database.
dataset : lambda
The creator of a dataset.
shuffle : bool, optional, default=False
Whether to shuffle the data.
num_chunks : int, optional, default=2048
num_chunks : int, optional, default=0
The number of chunks to split.
batch_size : int, optional, default=128
batch_size : int, optional, default=2
The size of a mini-batch.
prefetch : int, optional, default=5
The prefetch count.
"""
super(DataBatch, self).__init__()
# Init mpi
global_rank, local_rank, group_size = 0, 0, 1
if dragon.mpi.is_init():
group = dragon.mpi.is_parallel()
if group is not None: # DataParallel
global_rank = dragon.mpi.rank()
group_size = len(group)
for i, node in enumerate(group):
if global_rank == node:
local_rank = i
# Distributed settings
rank, group_size = 0, 1
process_group = dragon.distributed.get_default_process_group()
if process_group is not None and kwargs.get(
'phase', 'TRAIN') == 'TRAIN':
group_size = process_group.size
rank = dragon.distributed.get_rank(process_group)
kwargs['group_size'] = group_size
# Configuration
......@@ -71,6 +86,7 @@ class DataBatch(object):
self._num_transformers = kwargs.get('num_transformers', -1)
self._max_transformers = kwargs.get('max_transformers', 3)
self._num_fetchers = kwargs.get('num_fetchers', 1)
self.daemon = True
# Io-Aware Policy
if self._num_transformers == -1:
......@@ -81,66 +97,52 @@ class DataBatch(object):
self._num_transformers = min(
self._num_transformers, self._max_transformers)
# Init queues
self.Q1 = Queue(self._prefetch * self._num_readers * self._batch_size)
self.Q21 = Queue(self._prefetch * self._num_readers * self._batch_size)
self.Q22 = Queue(self._prefetch * self._num_readers * self._batch_size)
self.Q3 = Queue(self._prefetch * self._num_readers)
# Initialize queues
num_batches = self._prefetch * self._num_readers
self.Q1 = mp.Queue(num_batches * self._batch_size)
self.Q21 = mp.Queue(num_batches * self._batch_size)
self.Q22 = mp.Queue(num_batches * self._batch_size)
self.Q3 = mp.Queue(num_batches)
# Init readers
# Initialize readers
self._readers = []
for i in range(self._num_readers):
self._readers.append(DataReader(**kwargs))
self._readers[-1].q_out = self.Q1
for i in range(self._num_readers):
part_idx, num_parts = i, self._num_readers
num_parts *= group_size
part_idx += local_rank * self._num_readers
self._readers[i]._num_parts = num_parts
self._readers[i]._part_idx = part_idx
self._readers[i]._rng_seed += part_idx
part_idx += rank * self._num_readers
self._readers.append(dragon.io.DataReader(
num_parts=num_parts, part_idx=part_idx, **kwargs))
self._readers[i]._seed += part_idx
self._readers[i].q_out = self.Q1
self._readers[i].start()
time.sleep(0.1)
# Init transformers
# Initialize transformers
self._transformers = []
for i in range(self._num_transformers):
transformer = DataTransformer(**kwargs)
transformer._rng_seed += (i + local_rank * self._num_transformers)
transformer._rng_seed += (i + rank * self._num_transformers)
transformer.q_in = self.Q1
transformer.q1_out = self.Q21
transformer.q2_out = self.Q22
transformer.q1_out, transformer.q2_out = self.Q21, self.Q22
transformer.start()
self._transformers.append(transformer)
time.sleep(0.1)
# Init blob fetchers
self._fetchers = []
for i in range(self._num_fetchers):
fetcher = BlobFetcher(**kwargs)
fetcher.q1_in = self.Q21
fetcher.q2_in = self.Q22
fetcher.q_out = self.Q3
fetcher.start()
self._fetchers.append(fetcher)
time.sleep(0.1)
# Prevent to echo multiple nodes
if local_rank == 0:
self.echo()
# Initialize batch-producer
self.start()
# Register cleanup callbacks
def cleanup():
def terminate(processes):
for process in processes:
process.terminate()
process.join()
terminate(self._fetchers)
logger.info('Terminating BlobFetcher ......')
terminate([self])
logger.info('Terminate DataBatch.')
terminate(self._transformers)
logger.info('Terminating DataTransformer ......')
logger.info('Terminate DataTransformer.')
terminate(self._readers)
logger.info('Terminating DataReader......')
logger.info('Terminate DataReader.')
import atexit
atexit.register(cleanup)
......@@ -156,20 +158,27 @@ class DataBatch(object):
"""
return self.Q3.get()
def echo(self):
"""Print I/O Information.
Returns
-------
None
"""
print('---------------------------------------------------------')
print('BatchFetcher({} Threads), Using config:'.format(
self._num_readers + self._num_transformers + self._num_fetchers))
params = {'queue_size': self._prefetch,
'n_readers': self._num_readers,
'n_transformers': self._num_transformers,
'n_fetchers': self._num_fetchers}
pprint.pprint(params)
print('---------------------------------------------------------')
def run(self):
"""Start the process to produce batches."""
def produce(q_in):
processed_ims, ims_info, all_boxes = [], [], []
for image_index in range(cfg.TRAIN.IMS_PER_BATCH):
im, im_scale, gt_boxes = q_in.get()
processed_ims.append(im)
ims_info.append(list(im.shape[:2]) + [im_scale])
im_boxes = np.zeros((gt_boxes.shape[0], gt_boxes.shape[1] + 1), 'float32')
im_boxes[:, :gt_boxes.shape[1]], im_boxes[:, -1] = gt_boxes, image_index
all_boxes.append(im_boxes)
return {
'data': im_list_to_blob(processed_ims),
'ims_info': np.array(ims_info, dtype=np.float32),
'gt_boxes': np.concatenate(all_boxes, axis=0),
}
q1, q2 = self.Q21, self.Q22
while True:
if q1.qsize() >= cfg.TRAIN.IMS_PER_BATCH:
self.Q3.put(produce(q1))
elif q2.qsize() >= cfg.TRAIN.IMS_PER_BATCH:
self.Q3.put(produce(q2))
q1, q2 = q2, q1 # Sample two queues uniformly
lib/faster_rcnn/data/data_transformer.py → lib/faster_rcnn/data_transformer.py
......@@ -14,22 +14,13 @@ from __future__ import division
from __future__ import print_function
import multiprocessing
import numpy as np
import numpy.random as npr
try:
import cv2
except ImportError as e:
print('Failed to import cv2. Error: {0}'.format(str(e)))
try:
import PIL.Image
except ImportError as e:
print('Failed to import PIL. Error: {0}'.format(str(e)))
import cv2
import numpy as np
from lib.core.config import cfg
from lib.proto import anno_pb2 as pb
from lib.utils import logger
from lib.utils.blob import prep_im_for_blob
from lib.utils.boxes import flip_boxes
class DataTransformer(multiprocessing.Process):
......@@ -47,44 +38,45 @@ class DataTransformer(multiprocessing.Process):
def make_roi_dict(
self,
ann_datum,
example,
im_scale,
apply_flip=False,
offsets=None,
):
annotations = ann_datum.annotation
n_objects = 0
if not self._use_diff:
for ann in annotations:
if not ann.difficult:
for obj in example['object']:
if obj.get('difficult', 0) == 0:
n_objects += 1
else:
n_objects = len(annotations)
n_objects = len(example['object'])
roi_dict = {
'width': ann_datum.datum.width,
'height': ann_datum.datum.height,
'width': example['width'],
'height': example['height'],
'gt_classes': np.zeros((n_objects,), 'int32'),
'boxes': np.zeros((n_objects, 4), 'float32'),
}
# Filter the difficult instances
rec_idx = 0
for ann in annotations:
if not self._use_diff and ann.difficult:
object_idx = 0
for obj in example['object']:
if not self._use_diff and \
obj.get('difficult', 0) > 0:
continue
roi_dict['boxes'][rec_idx, :] = [
max(0, ann.x1),
max(0, ann.y1),
min(ann.x2, ann_datum.datum.width - 1),
min(ann.y2, ann_datum.datum.height - 1),
roi_dict['boxes'][object_idx, :] = [
max(0, obj['xmin']),
max(0, obj['ymin']),
min(obj['xmax'], example['width'] - 1),
min(obj['ymax'], example['height'] - 1),
]
roi_dict['gt_classes'][rec_idx] = self._class_to_ind[ann.name]
rec_idx += 1
roi_dict['gt_classes'][object_idx] = \
self._class_to_ind[obj['name']]
object_idx += 1
# Flip the boxes if necessary
if apply_flip:
roi_dict['boxes'] = _flip_boxes(
roi_dict['boxes'] = flip_boxes(
roi_dict['boxes'], roi_dict['width'])
# Scale the boxes to the detecting scale
......@@ -102,50 +94,34 @@ class DataTransformer(multiprocessing.Process):
return roi_dict
@classmethod
def get_image(cls, serialized):
datum = pb.AnnotatedDatum()
datum.ParseFromString(serialized)
datum = datum.datum
im = np.fromstring(datum.data, np.uint8)
return cv2.imdecode(im, -1) if datum.encoded is True else \
im.reshape((datum.height, datum.width, datum.channels))
def get_image(cls, example):
img = np.frombuffer(example['content'], np.uint8)
return cv2.imdecode(img, -1)
@classmethod
def get_annotations(cls, serialized):
datum = pb.AnnotatedDatum()
datum.ParseFromString(serialized)
filename = datum.filename
annotations = datum.annotation
def get_annotations(cls, example):
objects = []
for ix, ann in enumerate(annotations):
for ix, obj in enumerate(example['object']):
objects.append({
'name': ann.name,
'difficult': int(ann.difficult),
'bbox': [ann.x1, ann.y1, ann.x2, ann.y2],
'mask': ann.mask,
'name': obj['name'],
'difficult': obj.get('difficult', 0),
'bbox': [obj['xmin'], obj['ymin'], obj['xmax'], obj['ymax']],
})
return filename, objects
def get(self, serialized):
datum = pb.AnnotatedDatum()
datum.ParseFromString(serialized)
im_datum = datum.datum
im = np.fromstring(im_datum.data, np.uint8)
if im_datum.encoded is True:
im = cv2.imdecode(im, -1)
else:
h, w = im_datum.height, im_datum.width
im = im.reshape((h, w, im_datum.channels))
return example['id'], objects
def get(self, example):
img = np.frombuffer(example['content'], np.uint8)
img = cv2.imdecode(img, -1)
# Scale
scale_indices = npr.randint(len(cfg.TRAIN.SCALES))
scale_indices = np.random.randint(len(cfg.TRAIN.SCALES))
target_size = cfg.TRAIN.SCALES[scale_indices]
im, im_scale, jitter = prep_im_for_blob(im, target_size, cfg.TRAIN.MAX_SIZE)
im, im_scale, jitter = prep_im_for_blob(img, target_size, cfg.TRAIN.MAX_SIZE)
# Flip
apply_flip = False
if self._use_flipped:
if npr.randint(0, 2) > 0:
if np.random.randint(2) > 0:
im = im[:, ::-1, :]
apply_flip = True
......@@ -160,8 +136,8 @@ class DataTransformer(multiprocessing.Process):
# To a square (target_size, target_size)
im, offsets = _get_image_with_target_size([target_size] * 2, im)
# Datum -> RoIDict
roi_dict = self.make_roi_dict(datum, im_scale, apply_flip, offsets)
# Example -> RoIDict
roi_dict = self.make_roi_dict(example, im_scale, apply_flip, offsets)
# Post-Process for gt boxes
# Shape like: [num_objects, {x1, y1, x2, y2, cls}]
......@@ -171,29 +147,16 @@ class DataTransformer(multiprocessing.Process):
return im, im_scale, gt_boxes
def run(self):
npr.seed(self._rng_seed)
np.random.seed(self._rng_seed)
while True:
serialized = self.q_in.get()
data = self.get(serialized)
# Ensure that there should be at least 1 ground-truth
if len(data[2]) < 1:
continue
aspect_ratio = float(data[0].shape[0]) / data[0].shape[1]
if aspect_ratio > 1.0:
self.q1_out.put(data)
outputs = self.get(self.q_in.get())
if len(outputs[2]) < 1:
continue # Ignore the non-object image
aspect_ratio = float(outputs[0].shape[0]) / outputs[0].shape[1]
if aspect_ratio > 1.:
self.q1_out.put(outputs)
else:
self.q2_out.put(data)
def _flip_boxes(boxes, width):
flip_boxes = boxes.copy()
old_x1 = boxes[:, 0].copy()
old_x2 = boxes[:, 2].copy()
flip_boxes[:, 0] = width - old_x2 - 1
flip_boxes[:, 2] = width - old_x1 - 1
if not (flip_boxes[:, 2] >= flip_boxes[:, 0]).all():
logger.fatal('Encounter invalid coordinates after flipping boxes.')
return flip_boxes
self.q2_out.put(outputs)
def _get_image_with_target_size(target_size, img):
......
lib/faster_rcnn/generate_anchors.py
......@@ -13,6 +13,10 @@
#
# ------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
# Verify that we compute the same anchors as Shaoqing's matlab implementation:
......
lib/faster_rcnn/layers/__init__.py deleted 100644 → 0
# --------------------------------------------------------
# Mask R-CNN @ Detectron
# Copyright (c) 2017 SeetaTech
# Written by Ting Pan
# --------------------------------------------------------
\ No newline at end of file
lib/faster_rcnn/layers/data_layer.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 dragon.vm.torch as torch
from lib.core.config import cfg
from lib.datasets.factory import get_imdb
from lib.faster_rcnn.data.data_batch import DataBatch
class DataLayer(torch.nn.Module):
def __init__(self):
super(DataLayer, self).__init__()
database = get_imdb(cfg.TRAIN.DATABASE)
self.data_batch = DataBatch(**{
'source': database.source,
'classes': database.classes,
'shuffle': cfg.TRAIN.USE_SHUFFLE,
'num_chunks': 0, # Record-Wise Shuffle
'batch_size': cfg.TRAIN.IMS_PER_BATCH * 2,
})
def forward(self):
# Get an array blob from the Queue
outputs = self.data_batch.get()
# Zero-Copy the array to tensor
outputs['data'] = torch.from_numpy(outputs['data'])
return outputs
lib/faster_rcnn/test.py
......@@ -38,6 +38,7 @@ def im_detect(detector, raw_image):
blobs['ims_info'] = np.array([
list(blobs['data'].shape[1:3]) + [im_scale]
for im_scale in ims_scale], dtype=np.float32)
blobs['data'] = torch.from_numpy(blobs['data'])
# Do Forward
......@@ -129,8 +130,10 @@ def test_net(detector, server):
_t['misc'].toc()
print('\rim_detect: {:d}/{:d} {:.3f}s {:.3f}s'
.format(i + 1, num_images, _t['im_detect'].average_time,
_t['misc'].average_time), end='')
.format(i + 1, num_images,
_t['im_detect'].average_time,
_t['misc'].average_time),
end='')
print('\n>>>>>>>>>>>>>>>>>>> Evaluating <<<<<<<<<<<<<<<<<<<<')
......
lib/fpn/__init__.py
......@@ -13,6 +13,6 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from lib.fpn.layers.anchor_target_layer import AnchorTargetLayer
from lib.fpn.layers.proposal_layer import ProposalLayer
from lib.fpn.layers.proposal_target_layer import ProposalTargetLayer
from lib.fpn.anchor_target_layer import AnchorTargetLayer
from lib.fpn.proposal_layer import ProposalLayer
from lib.fpn.proposal_target_layer import ProposalTargetLayer
lib/fpn/layers/__init__.py deleted 100644 → 0
# --------------------------------------------------------
# Mask R-CNN @ Detectron
# Copyright (c) 2017 SeetaTech
# Written by Ting Pan
# --------------------------------------------------------
\ No newline at end of file
lib/modeling/detector.py
......@@ -35,11 +35,13 @@ class Detector(torch.nn.Module):
``lib.core.config`` for their hyper-parameters.
"""
def __init__(self):
super(Detector, self).__init__()
model = cfg.MODEL.TYPE
backbone = cfg.MODEL.BACKBONE.lower().split('.')
body, modules = backbone[0], backbone[1:]
self.recorder = None
# + Data Loader
self.data_layer = importlib.import_module(
......@@ -92,9 +94,14 @@ class Detector(torch.nn.Module):
Parameters
----------
inputs : dict or None
inputs : dict, optional
The inputs.
Returns
-------
dict
The outputs.
"""
# 0. Get the inputs
if inputs is None:
......@@ -161,7 +168,6 @@ class Detector(torch.nn.Module):
"""Optimize the graph for the inference.
It usually involves the removing of BN or Affine.
"""
##################################
# Merge Affine into Convolution #
......
lib/modeling/fast_rcnn.py
......@@ -54,7 +54,7 @@ class FastRCNN(torch.nn.Module):
'RoIAlign': torch.vision.ops.roi_align,
}[cfg.FRCNN.ROI_XFORM_METHOD]
self.cls_loss = torch.nn.CrossEntropyLoss(ignore_index=-1)
self.bbox_loss = torch.nn.SmoothL1Loss(beta=1., reduction='batch_size')
self.bbox_loss = torch.nn.SmoothL1Loss(reduction='batch_size')
# Compute spatial scales for multiple strides
roi_levels = [level for level in range(
cfg.FPN.ROI_MIN_LEVEL, cfg.FPN.ROI_MAX_LEVEL + 1)]
......@@ -130,25 +130,22 @@ class FastRCNN(torch.nn.Module):
# Compute rcnn logits
cls_score = self.cls_score(rcnn_output).float()
outputs = collections.OrderedDict({
'bbox_pred':
self.bbox_pred(rcnn_output).float(),
})
outputs = collections.OrderedDict([
('bbox_pred', self.bbox_pred(rcnn_output).float()),
])
if self.training:
# Compute rcnn losses
outputs.update(collections.OrderedDict({
'cls_loss': self.cls_loss(
cls_score,
self.rcnn_data['labels'],
),
'bbox_loss': self.bbox_loss(
outputs.update(collections.OrderedDict([
('cls_loss', self.cls_loss(
cls_score, self.rcnn_data['labels'])),
('bbox_loss', self.bbox_loss(
outputs['bbox_pred'],
self.rcnn_data['bbox_targets'],
self.rcnn_data['bbox_inside_weights'],
self.rcnn_data['bbox_outside_weights'],
),
}))
)),
]))
else:
# Return the rois to decode the refine boxes
if len(self.rcnn_data['rois']) > 1:
......
lib/modeling/fpn.py
......@@ -72,7 +72,7 @@ class FPN(torch.nn.Module):
def apply_on_retinanet(self, features):
fpn_input = self.C[-1](features[-1])
min_lvl, max_lvl = cfg.FPN.RPN_MIN_LEVEL, cfg.FPN.RPN_MAX_LEVEL
outputs = [self.P[HIGHEST_BACKBONE_LVL- min_lvl](fpn_input)]
outputs = [self.P[HIGHEST_BACKBONE_LVL - min_lvl](fpn_input)]
# Add extra convolutions for higher features
extra_input = features[-1]
for i in range(HIGHEST_BACKBONE_LVL + 1, max_lvl + 1):
......
lib/modeling/retinanet.py
......@@ -59,8 +59,7 @@ class RetinaNet(torch.nn.Module):
gamma=cfg.MODEL.FOCAL_LOSS_GAMMA,
)
self.bbox_loss = torch.nn.SmoothL1Loss(
beta=1. / 9., reduction='batch_size',
)
beta=.11, reduction='batch_size')
self.reset_parameters()
def reset_parameters(self):
......@@ -133,26 +132,22 @@ class RetinaNet(torch.nn.Module):
gt_boxes=gt_boxes,
ims_info=ims_info,
)
return collections.OrderedDict({
'cls_loss':
self.cls_loss(
cls_score,
self.retinanet_data['labels'],
),
'bbox_loss':
self.bbox_loss(
return collections.OrderedDict([
('cls_loss', self.cls_loss(
cls_score, self.retinanet_data['labels'])),
('bbox_loss', self.bbox_loss(
bbox_pred,
self.retinanet_data['bbox_targets'],
self.retinanet_data['bbox_inside_weights'],
self.retinanet_data['bbox_outside_weights'],
)
})
)),
])
def forward(self, *args, **kwargs):
cls_score, bbox_pred = self.compute_outputs(kwargs['features'])
cls_score, bbox_pred = cls_score.float(), bbox_pred.float()
outputs = collections.OrderedDict({'bbox_pred': bbox_pred})
outputs = collections.OrderedDict([('bbox_pred', bbox_pred)])
if self.training:
outputs.update(
......
lib/modeling/rpn.py
......@@ -44,14 +44,15 @@ class RPN(torch.nn.Module):
if len(cfg.RPN.STRIDES) > 1:
# RPN with multiple strides(i.e. FPN)
from lib.fpn.layers.anchor_target_layer import AnchorTargetLayer
from lib.fpn.anchor_target_layer import AnchorTargetLayer
else:
# RPN with single stride(i.e. C4)
from lib.faster_rcnn.layers.anchor_target_layer import AnchorTargetLayer
from lib.faster_rcnn.anchor_target_layer import AnchorTargetLayer
self.anchor_target_layer = AnchorTargetLayer()
self.cls_loss = torch.nn.BCEWithLogitsLoss()
self.bbox_loss = torch.nn.SmoothL1Loss(beta=1. / 9.)
self.bbox_loss = torch.nn.SmoothL1Loss(
beta=.11, reduction='batch_size')
self.reset_parameters()
def reset_parameters(self):
......@@ -120,26 +121,25 @@ class RPN(torch.nn.Module):
gt_boxes=gt_boxes,
ims_info=ims_info,
)
return collections.OrderedDict({
'rpn_cls_loss':
self.cls_loss(cls_score, self.rpn_data['labels']),
'rpn_bbox_loss':
self.bbox_loss(
return collections.OrderedDict([
('rpn_cls_loss', self.cls_loss(
cls_score, self.rpn_data['labels'])),
('rpn_bbox_loss', self.bbox_loss(
bbox_pred,
self.rpn_data['bbox_targets'],
self.rpn_data['bbox_inside_weights'],
self.rpn_data['bbox_outside_weights'],
)
})
)),
])
def forward(self, *args, **kwargs):
cls_score, bbox_pred = self.compute_outputs(kwargs['features'])
cls_score, bbox_pred = cls_score.float(), bbox_pred.float()
outputs = collections.OrderedDict({
'rpn_cls_score': cls_score,
'rpn_bbox_pred': bbox_pred,
})
outputs = collections.OrderedDict([
('rpn_cls_score', cls_score),
('rpn_bbox_pred', bbox_pred),
])
if self.training:
outputs.update(
......
lib/modeling/ssd.py
......@@ -136,32 +136,29 @@ class SSD(torch.nn.Module):
gt_boxes=gt_boxes,
)
)
return collections.OrderedDict({
return collections.OrderedDict([
# A compensating factor of 4.0 is used
# As we normalize both the pos and neg samples
'cls_loss':
self.cls_loss(
('cls_loss', self.cls_loss(
cls_score.view(-1, cfg.MODEL.NUM_CLASSES),
self.ssd_data['labels']
) * 4.,
'bbox_loss':
self.bbox_loss(
self.ssd_data['labels']) * 4.),
('bbox_loss', self.bbox_loss(
bbox_pred,
self.ssd_data['bbox_targets'],
self.ssd_data['bbox_inside_weights'],
self.ssd_data['bbox_outside_weights'],
)
})
)),
])
def forward(self, *args, **kwargs):
prior_boxes = self.prior_box_layer(kwargs['features'])
cls_score, bbox_pred = self.compute_outputs(kwargs['features'])
cls_score, bbox_pred = cls_score.float(), bbox_pred.float()
outputs = collections.OrderedDict({
'prior_boxes': prior_boxes,
'bbox_pred': bbox_pred,
})
outputs = collections.OrderedDict([
('bbox_pred', bbox_pred),
('prior_boxes', prior_boxes),
])
if self.training:
outputs.update(
......
lib/proto/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
\ No newline at end of file
lib/proto/anno.proto deleted 100644 → 0
syntax = "proto2";
message Datum {
optional int32 channels = 1;
optional int32 height = 2;
optional int32 width = 3;
optional bytes data = 4;
optional int32 label = 5;
repeated float float_data = 6;
optional bool encoded = 7 [default = false];
repeated int32 labels = 8;
}
message Annotation {
optional float x1 = 1;
optional float y1 = 2;
optional float x2 = 3;
optional float y2 = 4;
optional string name = 5;
optional bool difficult = 6 [default = false];
optional string mask = 7;
}
message AnnotatedDatum {
optional Datum datum = 1;
optional string filename = 2;
repeated Annotation annotation = 3;
}
lib/retinanet/__init__.py
......@@ -13,5 +13,5 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from lib.faster_rcnn.layers.data_layer import DataLayer
from lib.retinanet.layers.anchor_target_layer import AnchorTargetLayer
from lib.faster_rcnn.data_layer import DataLayer
from lib.retinanet.anchor_target_layer import AnchorTargetLayer
lib/retinanet/layers/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
lib/retinanet/test.py
......@@ -172,8 +172,10 @@ def test_net(net, server):
print('\rim_detect: {:d}/{:d} {:.3f}s {:.3f}s'
.format(batch_idx + cfg.TEST.IMS_PER_BATCH,
num_images, _t['im_detect'].average_time,
_t['misc'].average_time), end='')
num_images,
_t['im_detect'].average_time,
_t['misc'].average_time),
end='')
print('\n>>>>>>>>>>>>>>>>>>> Evaluating <<<<<<<<<<<<<<<<<<<<')
......
lib/ssd/__init__.py
......@@ -13,8 +13,8 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from lib.ssd.layers.data_layer import DataLayer
from lib.ssd.layers.hard_mining_layer import HardMiningLayer
from lib.ssd.layers.multibox_layer import MultiBoxMatchLayer
from lib.ssd.layers.multibox_layer import MultiBoxTargetLayer
from lib.ssd.layers.priorbox_layer import PriorBoxLayer
from lib.ssd.data_layer import DataLayer
from lib.ssd.hard_mining_layer import HardMiningLayer
from lib.ssd.multibox_layer import MultiBoxMatchLayer
from lib.ssd.multibox_layer import MultiBoxTargetLayer
from lib.ssd.priorbox_layer import PriorBoxLayer
lib/ssd/data/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
\ No newline at end of file
lib/ssd/data/blob_fetcher.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 multiprocessing
import numpy as np
from lib.core.config import cfg
class BlobFetcher(multiprocessing.Process):
def __init__(self, **kwargs):
super(BlobFetcher, self).__init__()
self._img_blob_size = (
cfg.TRAIN.IMS_PER_BATCH,
cfg.SSD.RESIZE.HEIGHT,
cfg.SSD.RESIZE.WIDTH, 3,
)
self.q_in = self.q_out = None
self.daemon = True
def get(self):
img_blob, boxes_blob = np.zeros(self._img_blob_size, 'uint8'), []
for i in range(cfg.TRAIN.IMS_PER_BATCH):
img_blob[i], gt_boxes = self.q_in.get()
# Pack the boxes by adding the index of images
boxes = np.zeros((gt_boxes.shape[0], gt_boxes.shape[1] + 1), np.float32)
boxes[:, :gt_boxes.shape[1]] = gt_boxes
boxes[:, -1] = i
boxes_blob.append(boxes)
return {
'data': img_blob,
'gt_boxes': np.concatenate(boxes_blob, 0),
}
def run(self):
while True:
self.q_out.put(self.get())
lib/ssd/data/data_batch.py → lib/ssd/data_layer.py
......@@ -13,54 +13,69 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from multiprocessing import Queue
import multiprocessing as mp
import time
import dragon
import pprint
import dragon.vm.torch as torch
import numpy as np
from lib.faster_rcnn.data.data_reader import DataReader
from lib.ssd.data.data_transformer import DataTransformer
from lib.ssd.data.blob_fetcher import BlobFetcher
from lib.core.config import cfg
from lib.datasets.factory import get_imdb
from lib.ssd.data_transformer import DataTransformer
from lib.utils import logger
class DataBatch(object):
"""DataBatch aims to prefetch data by ``Triple-Buffering``.
class DataLayer(torch.nn.Module):
"""Generate a mini-batch of data."""
It takes full advantages of the Process/Thread of Python,
def __init__(self):
super(DataLayer, self).__init__()
database = get_imdb(cfg.TRAIN.DATABASE)
self.data_batch = DataBatch(**{
'dataset': lambda: dragon.io.SeetaRecordDataset(database.source),
'classes': database.classes,
'shuffle': cfg.TRAIN.USE_SHUFFLE,
'num_chunks': cfg.TRAIN.NUM_SHUFFLE_CHUNKS,
'batch_size': cfg.TRAIN.IMS_PER_BATCH * 2,
})
which provides remarkable I/O speed up for scalable distributed training.
def forward(self):
# Get an array blob from the Queue
outputs = self.data_batch.get()
# Zero-Copy the array to tensor
outputs['data'] = torch.from_numpy(outputs['data'])
return outputs
class DataBatch(mp.Process):
"""Prefetch the batch of data."""
"""
def __init__(self, **kwargs):
"""Construct a ``DataBatch``.
Parameters
----------
source : str
The path of database.
dataset : lambda
The creator of a dataset.
shuffle : bool, optional, default=False
Whether to shuffle the data.
num_chunks : int, optional, default=2048
num_chunks : int, optional, default=0
The number of chunks to split.
batch_size : int, optional, default=128
batch_size : int, optional, default=32
The size of a mini-batch.
prefetch : int, optional, default=5
The prefetch count.
"""
super(DataBatch, self).__init__()
# Init mpi
global_rank, local_rank, group_size = 0, 0, 1
if dragon.mpi.is_init():
group = dragon.mpi.is_parallel()
if group is not None: # DataParallel
global_rank = dragon.mpi.rank()
group_size = len(group)
for i, node in enumerate(group):
if global_rank == node:
local_rank = i
# Distributed settings
rank, group_size = 0, 1
process_group = dragon.distributed.get_default_process_group()
if process_group is not None and kwargs.get(
'phase', 'TRAIN') == 'TRAIN':
group_size = process_group.size
rank = dragon.distributed.get_rank(process_group)
kwargs['group_size'] = group_size
# Configuration
......@@ -77,63 +92,50 @@ class DataBatch(object):
self._num_transformers = min(
self._num_transformers, self._max_transformers)
# Init queues
self.Q1 = Queue(self._prefetch * self._num_readers * self._batch_size)
self.Q2 = Queue(self._prefetch * self._num_readers * self._batch_size)
self.Q3 = Queue(self._prefetch * self._num_readers)
# Initialize queues
num_batches = self._prefetch * self._num_readers
self.Q1 = mp.Queue(num_batches * self._batch_size)
self.Q2 = mp.Queue(num_batches * self._batch_size)
self.Q3 = mp.Queue(num_batches)
# Init readers
# Initialize readers
self._readers = []
for i in range(self._num_readers):
self._readers.append(DataReader(**kwargs))
self._readers[-1].q_out = self.Q1
for i in range(self._num_readers):
part_idx, num_parts = i, self._num_readers
num_parts *= group_size
part_idx += local_rank * self._num_readers
self._readers[i]._num_parts = num_parts
self._readers[i]._part_idx = part_idx
self._readers[i]._rng_seed += part_idx
part_idx += rank * self._num_readers
self._readers.append(dragon.io.DataReader(
num_parts=num_parts, part_idx=part_idx, **kwargs))
self._readers[i]._seed += part_idx
self._readers[i].q_out = self.Q1
self._readers[i].start()
time.sleep(0.1)
# Init transformers
# Initialize transformers
self._transformers = []
for i in range(self._num_transformers):
transformer = DataTransformer(**kwargs)
transformer._rng_seed += (i + local_rank * self._num_transformers)
transformer.q_in = self.Q1
transformer.q_out = self.Q2
transformer._rng_seed += (i + rank * self._num_transformers)
transformer.q_in, transformer.q_out = self.Q1, self.Q2
transformer.start()
self._transformers.append(transformer)
time.sleep(0.1)
# Init blob fetchers
self._fetchers = []
for i in range(self._num_fetchers):
fetcher = BlobFetcher(**kwargs)
fetcher.q_in = self.Q2
fetcher.q_out = self.Q3
fetcher.start()
self._fetchers.append(fetcher)
time.sleep(0.1)
# Prevent to echo multiple nodes
if local_rank == 0:
self.echo()
# Initialize batch-producer
self.start()
# Register cleanup callbacks
def cleanup():
def terminate(processes):
for process in processes:
process.terminate()
process.join()
terminate(self._fetchers)
logger.info('Terminating BlobFetcher ......')
terminate([self])
logger.info('Terminate DataBatch.')
terminate(self._transformers)
logger.info('Terminating DataTransformer ......')
logger.info('Terminate DataTransformer.')
terminate(self._readers)
logger.info('Terminating DataReader......')
logger.info('Terminate DataReader.')
import atexit
atexit.register(cleanup)
......@@ -149,14 +151,24 @@ class DataBatch(object):
"""
return self.Q3.get()
def echo(self):
"""Print I/O Information."""
print('---------------------------------------------------------')
print('BatchFetcher({} Threads), Using config:'.format(
self._num_readers + self._num_transformers + self._num_fetchers))
params = {'queue_size': self._prefetch,
'n_readers': self._num_readers,
'n_transformers': self._num_transformers,
'n_fetchers': self._num_fetchers}
pprint.pprint(params)
print('---------------------------------------------------------')
def run(self):
"""Start the process to produce batches."""
image_batch_shape = (
cfg.TRAIN.IMS_PER_BATCH,
cfg.SSD.RESIZE.HEIGHT,
cfg.SSD.RESIZE.WIDTH, 3,
)
while True:
boxes_to_pack = []
image_batch = np.zeros(image_batch_shape, 'uint8')
for image_index in range(cfg.TRAIN.IMS_PER_BATCH):
image_batch[image_index], gt_boxes = self.Q2.get()
boxes = np.zeros((gt_boxes.shape[0], gt_boxes.shape[1] + 1), 'float32')
boxes[:, :gt_boxes.shape[1]], boxes[:, -1] = gt_boxes, image_index
boxes_to_pack.append(boxes)
self.Q3.put({
'data': image_batch,
'gt_boxes': np.concatenate(boxes_to_pack),
})
lib/ssd/data/data_transformer.py → lib/ssd/data_transformer.py
......@@ -13,14 +13,14 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import cv2
import multiprocessing
import cv2
import numpy as np
from lib.core.config import cfg
from lib.proto import anno_pb2 as pb
from lib.ssd.data import transforms
from lib.utils import logger
from lib.ssd import transforms
from lib.utils.boxes import flip_boxes
class DataTransformer(multiprocessing.Process):
......@@ -41,38 +41,41 @@ class DataTransformer(multiprocessing.Process):
self.q_in = self.q_out = None
self.daemon = True
def make_roi_dict(self, ann_datum, flip=False):
annotations = ann_datum.annotation
def make_roi_dict(self, example, flip=False):
n_objects = 0
if not self._use_diff:
for ann in annotations:
if not ann.difficult: n_objects += 1
else: n_objects = len(annotations)
for obj in example['object']:
if obj.get('difficult', 0) == 0:
n_objects += 1
else:
n_objects = len(example['object'])
roi_dict = {
'width': ann_datum.datum.width,
'height': ann_datum.datum.height,
'gt_classes': np.zeros((n_objects,), dtype=np.int32),
'boxes': np.zeros((n_objects, 4), dtype=np.float32),
'normalized_boxes': np.zeros((n_objects, 4), dtype=np.float32),
'width': example['width'],
'height': example['height'],
'gt_classes': np.zeros((n_objects,), 'int32'),
'boxes': np.zeros((n_objects, 4), 'float32'),
'normalized_boxes': np.zeros((n_objects, 4), 'float32'),
}
rec_idx = 0
for ann in annotations:
if not self._use_diff and ann.difficult:
# Filter the difficult instances
object_idx = 0
for obj in example['object']:
if not self._use_diff and \
obj.get('difficult', 0) > 0:
continue
roi_dict['boxes'][rec_idx, :] = [
max(0, ann.x1),
max(0, ann.y1),
min(ann.x2, ann_datum.datum.width - 1),
min(ann.y2, ann_datum.datum.height - 1),
roi_dict['boxes'][object_idx, :] = [
max(0, obj['xmin']),
max(0, obj['ymin']),
min(obj['xmax'], example['width'] - 1),
min(obj['ymax'], example['height'] - 1),
]
roi_dict['gt_classes'][rec_idx] = \
self._class_to_ind[ann.name]
rec_idx += 1
roi_dict['gt_classes'][object_idx] = \
self._class_to_ind[obj['name']]
object_idx += 1
if flip:
roi_dict['boxes'] = _flip_boxes(
roi_dict['boxes'] = flip_boxes(
roi_dict['boxes'], roi_dict['width'])
roi_dict['boxes'][:, 0::2] /= roi_dict['width']
......@@ -80,26 +83,19 @@ class DataTransformer(multiprocessing.Process):
return roi_dict
def get(self, serialized):
ann_datum = pb.AnnotatedDatum()
ann_datum.ParseFromString(serialized)
img_datum = ann_datum.datum
img = np.fromstring(img_datum.data, np.uint8)
if img_datum.encoded is True:
def get(self, example):
img = np.frombuffer(example['content'], np.uint8)
img = cv2.imdecode(img, -1)
else:
h, w = img_datum.height, img_datum.width
img = img.reshape((h, w, img_datum.channels))
# Flip
flip = False
if self._mirror:
if np.random.randint(0, 2) > 0:
if np.random.randint(2) > 0:
img = img[:, ::-1, :]
flip = True
# Datum -> RoIDB
roi_dict = self.make_roi_dict(ann_datum, flip)
# Example -> RoIDict
roi_dict = self.make_roi_dict(example, flip)
# Post-Process for gt boxes
# Shape like: [num_objects, {x1, y1, x2, y2, cls}]
......@@ -120,19 +116,7 @@ class DataTransformer(multiprocessing.Process):
def run(self):
np.random.seed(self._rng_seed)
while True:
serialized = self.q_in.get()
im, gt_boxes = self.get(serialized)
if len(gt_boxes) < 1:
continue
self.q_out.put((im, gt_boxes))
def _flip_boxes(boxes, width):
flip_boxes = boxes.copy()
old_x1 = boxes[:, 0].copy()
old_x2 = boxes[:, 2].copy()
flip_boxes[:, 0] = width - old_x2 - 1
flip_boxes[:, 2] = width - old_x1 - 1
if not (flip_boxes[:, 2] >= flip_boxes[:, 0]).all():
logger.fatal('Encounter invalid coordinates after flipping boxes.')
return flip_boxes
outputs = self.get(self.q_in.get())
if len(outputs[1]) < 1:
continue # Ignore the non-object image
self.q_out.put(outputs)
lib/ssd/layers/__init__.py deleted 100644 → 0
# ------------------------------------------------------------
# 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>
#
# ------------------------------------------------------------
\ No newline at end of file
lib/ssd/layers/data_layer.py deleted 100644 → 0
# ------------------------------------------------------------
# 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 dragon.vm.torch as torch
from lib.core.config import cfg
from lib.datasets.factory import get_imdb
from lib.ssd.data.data_batch import DataBatch
class DataLayer(torch.nn.Module):
def __init__(self):
super(DataLayer, self).__init__()
database = get_imdb(cfg.TRAIN.DATABASE)
self.data_batch = DataBatch(**{
'source': database.source,
'classes': database.classes,
'shuffle': cfg.TRAIN.USE_SHUFFLE,
'num_chunks': 2048, # Chunk-Wise Shuffle
'batch_size': cfg.TRAIN.IMS_PER_BATCH * 2,
})
def forward(self):
# Get an array blob from the Queue
outputs = self.data_batch.get()
# Zero-Copy the array to tensor
outputs['data'] = torch.from_numpy(outputs['data'])
return outputs
lib/ssd/test.py
......@@ -152,8 +152,10 @@ def test_net(net, server):
print('\rim_detect: {:d}/{:d} {:.3f}s {:.3f}s'
.format(batch_idx + cfg.TEST.IMS_PER_BATCH,
num_images, _t['im_detect'].average_time,
_t['misc'].average_time), end='')
num_images,
_t['im_detect'].average_time,
_t['misc'].average_time),
end='')
print('\n>>>>>>>>>>>>>>>>>>> Evaluating <<<<<<<<<<<<<<<<<<<<')
......
lib/ssd/data/transforms_test.py → lib/ssd/transforms_test.py
......@@ -19,7 +19,7 @@ sys.path.append('../../')
import cv2
import numpy as np
from lib.ssd.data import transforms
from lib.ssd import transforms
if __name__ == '__main__':
......
lib/utils/boxes.py
......@@ -201,6 +201,16 @@ def expand_boxes(boxes, scale):
return boxes_exp
def flip_boxes(boxes, width):
"""Flip the boxes horizontally."""
flip_boxes = boxes.copy()
old_x1 = boxes[:, 0].copy()
old_x2 = boxes[:, 2].copy()
flip_boxes[:, 0] = width - old_x2 - 1
flip_boxes[:, 2] = width - old_x1 - 1
return flip_boxes
def filter_boxes(boxes, min_size):
"""Remove all boxes with any side smaller than min size."""
ws = boxes[:, 2] - boxes[:, 0] + 1
......
tools/mpi_train.py
......@@ -62,22 +62,20 @@ if __name__ == '__main__':
if checkpoint is not None:
cfg.TRAIN.WEIGHTS = checkpoint
# Setup MPI
if cfg.NUM_GPUS != dragon.mpi.size():
# Setup the distributed environment
world_rank = dragon.distributed.get_rank()
world_size = dragon.distributed.get_world_size()
if cfg.NUM_GPUS != world_size:
raise ValueError(
'Excepted {} mpi nodes, but got {}.'
.format(len(args.gpus), dragon.mpi.size())
'Excepted staring of {} processes, got {}.'
.format(cfg.NUM_GPUS, world_size)
)
GPUs = [i for i in range(cfg.NUM_GPUS)]
cfg.GPU_ID = GPUs[dragon.mpi.rank()]
dragon.mpi.add_parallel_group([i for i in range(cfg.NUM_GPUS)])
dragon.mpi.set_parallel_mode('NCCL' if cfg.USE_NCCL else 'MPI')
logger.set_root_logger(world_rank == 0)
# Setup logger
if dragon.mpi.rank() != 0:
logger.set_root_logger(False)
# Select the GPU depending on the rank of process
cfg.GPU_ID = [i for i in range(cfg.NUM_GPUS)][world_rank]
# Fix the random seeds (numpy and dragon) for reproducibility
# Fix the random seed for reproducibility
numpy.random.seed(cfg.RNG_SEED)
dragon.config.set_random_seed(cfg.RNG_SEED)
......@@ -89,7 +87,8 @@ if __name__ == '__main__':
# Ready to train the network
logger.info('Output will be saved to `{:s}`'
.format(coordinator.checkpoints_dir()))
with dragon.distributed.new_group(
ranks=[i for i in range(cfg.NUM_GPUS)],
backend='NCCL' if cfg.USE_NCCL else 'MPI',
verbose=True).as_default():
train_net(coordinator, start_iter)
# Finalize mpi
dragon.mpi.finalize()
tools/train.py
......@@ -82,7 +82,7 @@ if __name__ == '__main__':
if checkpoint is not None:
cfg.TRAIN.WEIGHTS = checkpoint
# Fix the random seeds (numpy and dragon) for reproducibility
# Fix the random seed for reproducibility
numpy.random.seed(cfg.RNG_SEED)
dragon.config.set_random_seed(cfg.RNG_SEED)
......
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!