class DataConfig(Config):
"""
Configuration for training on the Box dataset.
Derives from the base Config class and overrides values specific
to the Box dataset.
"""
# Give the configuration a recognizable name
NAME = dataset_name
# settle the gpu count and images per gpu by config in libs.
GPU_COUNT = len(cfgs.train_gpus.split(','))
IMAGES_PER_GPU = cfgs.images_per_gpu
# Number of classes (including background)
NUM_CLASSES = 1 + len(label_name_dict.keys()) # including the background
# Use small images for faster training. Set the limits of the small side
# the large side, and that determines the image shape.
IMAGE_MIN_DIM = cfgs.image_min_dim
IMAGE_MAX_DIM = cfgs.image_max_dim
# Use smaller anchors because our image and objects are small
RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128) # anchor side in pixels
# Reduce training ROIs per image because the images are small and have
# few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
TRAIN_ROIS_PER_IMAGE = 32
# Use a small epoch since the data is simple
STEPS_PER_EPOCH = 100
# use small validation steps since the epoch is small
VALIDATION_STEPS = 5
def load_data(self, image_dir, mask_dir, yaml_dir):
"""
load box
:param image_dir: image dir
:param mask_dir: mask dir
:param yaml_dir: yaml dir
:return:
"""
# Add classes
for label in label_name_dict.keys():
name = label_name_dict[label]
self.add_class(dataset_name, label, name)
# get all rgb files
image_files = [
file for file in os.listdir(image_dir)
if file.endswith('png') or file.endswith('jpg')
]
image_id = 0
for image_file in image_files:
yaml_file = image_file.replace('png', 'yaml')
mask_file = image_file
yaml_path = os.path.join(yaml_dir, yaml_file)
mask_path = os.path.join(mask_dir, mask_file)
image_path = os.path.join(image_dir, image_file)
img = cv2.imread(image_path)
height, width = img.shape[0], img.shape[1]
if not (os.path.exists(yaml_path) and os.path.exists(mask_path)):
# when yaml or mask file do not exist, continue without adding the images
continue
self.add_image(source=dataset_name,
image_id=image_id,
path=image_path,
width=width,
height=height,
mask_path=mask_path,
yaml_path=yaml_path)
image_id += 1