|
|
|
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
|
|
#
|
|
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
# you may not use this file except in compliance with the License.
|
|
|
|
# You may obtain a copy of the License at
|
|
|
|
#
|
|
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
#
|
|
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
# See the License for the specific language governing permissions and
|
|
|
|
# limitations under the License.
|
|
|
|
|
|
|
|
from __future__ import absolute_import
|
|
|
|
import copy
|
|
|
|
import os
|
|
|
|
import os.path as osp
|
|
|
|
import random
|
|
|
|
import re
|
|
|
|
from collections import OrderedDict
|
|
|
|
import xml.etree.ElementTree as ET
|
|
|
|
|
|
|
|
import numpy as np
|
|
|
|
from paddle.io import Dataset
|
|
|
|
|
|
|
|
from paddlers.utils import logging, get_num_workers, get_encoding, path_normalization, is_pic
|
|
|
|
from paddlers.transforms import ImgDecoder, MixupImage
|
|
|
|
from paddlers.tools import YOLOAnchorCluster
|
|
|
|
|
|
|
|
|
|
|
|
class VOCDetection(Dataset):
|
|
|
|
"""读取PascalVOC格式的检测数据集,并对样本进行相应的处理。
|
|
|
|
|
|
|
|
Args:
|
|
|
|
data_dir (str): 数据集所在的目录路径。
|
|
|
|
file_list (str): 描述数据集图片文件和对应标注文件的文件路径(文本内每行路径为相对data_dir的相对路)。
|
|
|
|
label_list (str): 描述数据集包含的类别信息文件路径。
|
|
|
|
transforms (paddlers.det.transforms): 数据集中每个样本的预处理/增强算子。
|
|
|
|
num_workers (int|str): 数据集中样本在预处理过程中的线程或进程数。默认为'auto'。当设为'auto'时,根据
|
|
|
|
系统的实际CPU核数设置`num_workers`: 如果CPU核数的一半大于8,则`num_workers`为8,否则为CPU核数的
|
|
|
|
一半。
|
|
|
|
shuffle (bool): 是否需要对数据集中样本打乱顺序。默认为False。
|
|
|
|
allow_empty (bool): 是否加载负样本。默认为False。
|
|
|
|
empty_ratio (float): 用于指定负样本占总样本数的比例。如果小于0或大于等于1,则保留全部的负样本。默认为1。
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self,
|
|
|
|
data_dir,
|
|
|
|
file_list,
|
|
|
|
label_list,
|
|
|
|
transforms=None,
|
|
|
|
num_workers='auto',
|
|
|
|
shuffle=False,
|
|
|
|
allow_empty=False,
|
|
|
|
empty_ratio=1.):
|
|
|
|
# matplotlib.use() must be called *before* pylab, matplotlib.pyplot,
|
|
|
|
# or matplotlib.backends is imported for the first time
|
|
|
|
# pycocotools import matplotlib
|
|
|
|
import matplotlib
|
|
|
|
matplotlib.use('Agg')
|
|
|
|
from pycocotools.coco import COCO
|
|
|
|
super(VOCDetection, self).__init__()
|
|
|
|
self.data_dir = data_dir
|
|
|
|
self.data_fields = None
|
|
|
|
self.transforms = copy.deepcopy(transforms)
|
|
|
|
self.num_max_boxes = 50
|
|
|
|
|
|
|
|
self.use_mix = False
|
|
|
|
if self.transforms is not None:
|
|
|
|
for op in self.transforms.transforms:
|
|
|
|
if isinstance(op, MixupImage):
|
|
|
|
self.mixup_op = copy.deepcopy(op)
|
|
|
|
self.use_mix = True
|
|
|
|
self.num_max_boxes *= 2
|
|
|
|
break
|
|
|
|
|
|
|
|
self.batch_transforms = None
|
|
|
|
self.num_workers = get_num_workers(num_workers)
|
|
|
|
self.shuffle = shuffle
|
|
|
|
self.allow_empty = allow_empty
|
|
|
|
self.empty_ratio = empty_ratio
|
|
|
|
self.file_list = list()
|
|
|
|
neg_file_list = list()
|
|
|
|
self.labels = list()
|
|
|
|
|
|
|
|
annotations = dict()
|
|
|
|
annotations['images'] = list()
|
|
|
|
annotations['categories'] = list()
|
|
|
|
annotations['annotations'] = list()
|
|
|
|
|
|
|
|
cname2cid = OrderedDict()
|
|
|
|
label_id = 0
|
|
|
|
with open(label_list, 'r', encoding=get_encoding(label_list)) as f:
|
|
|
|
for line in f.readlines():
|
|
|
|
cname2cid[line.strip()] = label_id
|
|
|
|
label_id += 1
|
|
|
|
self.labels.append(line.strip())
|
|
|
|
logging.info("Starting to read file list from dataset...")
|
|
|
|
for k, v in cname2cid.items():
|
|
|
|
annotations['categories'].append({
|
|
|
|
'supercategory': 'component',
|
|
|
|
'id': v + 1,
|
|
|
|
'name': k
|
|
|
|
})
|
|
|
|
ct = 0
|
|
|
|
ann_ct = 0
|
|
|
|
with open(file_list, 'r', encoding=get_encoding(file_list)) as f:
|
|
|
|
while True:
|
|
|
|
line = f.readline()
|
|
|
|
if not line:
|
|
|
|
break
|
|
|
|
if len(line.strip().split()) > 2:
|
|
|
|
raise Exception("A space is defined as the separator, "
|
|
|
|
"but it exists in image or label name {}."
|
|
|
|
.format(line))
|
|
|
|
img_file, xml_file = [
|
|
|
|
osp.join(data_dir, x) for x in line.strip().split()[:2]
|
|
|
|
]
|
|
|
|
img_file = path_normalization(img_file)
|
|
|
|
xml_file = path_normalization(xml_file)
|
|
|
|
if not is_pic(img_file):
|
|
|
|
continue
|
|
|
|
if not osp.isfile(xml_file):
|
|
|
|
continue
|
|
|
|
if not osp.exists(img_file):
|
|
|
|
logging.warning('The image file {} does not exist!'.format(
|
|
|
|
img_file))
|
|
|
|
continue
|
|
|
|
if not osp.exists(xml_file):
|
|
|
|
logging.warning('The annotation file {} does not exist!'.
|
|
|
|
format(xml_file))
|
|
|
|
continue
|
|
|
|
tree = ET.parse(xml_file)
|
|
|
|
if tree.find('id') is None:
|
|
|
|
im_id = np.asarray([ct])
|
|
|
|
else:
|
|
|
|
ct = int(tree.find('id').text)
|
|
|
|
im_id = np.asarray([int(tree.find('id').text)])
|
|
|
|
pattern = re.compile('<size>', re.IGNORECASE)
|
|
|
|
size_tag = pattern.findall(str(ET.tostringlist(tree.getroot())))
|
|
|
|
if len(size_tag) > 0:
|
|
|
|
size_tag = size_tag[0][1:-1]
|
|
|
|
size_element = tree.find(size_tag)
|
|
|
|
pattern = re.compile('<width>', re.IGNORECASE)
|
|
|
|
width_tag = pattern.findall(
|
|
|
|
str(ET.tostringlist(size_element)))[0][1:-1]
|
|
|
|
im_w = float(size_element.find(width_tag).text)
|
|
|
|
pattern = re.compile('<height>', re.IGNORECASE)
|
|
|
|
height_tag = pattern.findall(
|
|
|
|
str(ET.tostringlist(size_element)))[0][1:-1]
|
|
|
|
im_h = float(size_element.find(height_tag).text)
|
|
|
|
else:
|
|
|
|
im_w = 0
|
|
|
|
im_h = 0
|
|
|
|
|
|
|
|
pattern = re.compile('<object>', re.IGNORECASE)
|
|
|
|
obj_match = pattern.findall(
|
|
|
|
str(ET.tostringlist(tree.getroot())))
|
|
|
|
if len(obj_match) > 0:
|
|
|
|
obj_tag = obj_match[0][1:-1]
|
|
|
|
objs = tree.findall(obj_tag)
|
|
|
|
else:
|
|
|
|
objs = list()
|
|
|
|
|
|
|
|
num_bbox, i = len(objs), 0
|
|
|
|
gt_bbox = np.zeros((num_bbox, 4), dtype=np.float32)
|
|
|
|
gt_class = np.zeros((num_bbox, 1), dtype=np.int32)
|
|
|
|
gt_score = np.zeros((num_bbox, 1), dtype=np.float32)
|
|
|
|
is_crowd = np.zeros((num_bbox, 1), dtype=np.int32)
|
|
|
|
difficult = np.zeros((num_bbox, 1), dtype=np.int32)
|
|
|
|
for obj in objs:
|
|
|
|
pattern = re.compile('<name>', re.IGNORECASE)
|
|
|
|
name_tag = pattern.findall(str(ET.tostringlist(obj)))[0][1:
|
|
|
|
-1]
|
|
|
|
cname = obj.find(name_tag).text.strip()
|
|
|
|
pattern = re.compile('<difficult>', re.IGNORECASE)
|
|
|
|
diff_tag = pattern.findall(str(ET.tostringlist(obj)))
|
|
|
|
if len(diff_tag) == 0:
|
|
|
|
_difficult = 0
|
|
|
|
else:
|
|
|
|
diff_tag = diff_tag[0][1:-1]
|
|
|
|
try:
|
|
|
|
_difficult = int(obj.find(diff_tag).text)
|
|
|
|
except Exception:
|
|
|
|
_difficult = 0
|
|
|
|
pattern = re.compile('<bndbox>', re.IGNORECASE)
|
|
|
|
box_tag = pattern.findall(str(ET.tostringlist(obj)))
|
|
|
|
if len(box_tag) == 0:
|
|
|
|
logging.warning(
|
|
|
|
"There's no field '<bndbox>' in one of object, "
|
|
|
|
"so this object will be ignored. xml file: {}".
|
|
|
|
format(xml_file))
|
|
|
|
continue
|
|
|
|
box_tag = box_tag[0][1:-1]
|
|
|
|
box_element = obj.find(box_tag)
|
|
|
|
pattern = re.compile('<xmin>', re.IGNORECASE)
|
|
|
|
xmin_tag = pattern.findall(
|
|
|
|
str(ET.tostringlist(box_element)))[0][1:-1]
|
|
|
|
x1 = float(box_element.find(xmin_tag).text)
|
|
|
|
pattern = re.compile('<ymin>', re.IGNORECASE)
|
|
|
|
ymin_tag = pattern.findall(
|
|
|
|
str(ET.tostringlist(box_element)))[0][1:-1]
|
|
|
|
y1 = float(box_element.find(ymin_tag).text)
|
|
|
|
pattern = re.compile('<xmax>', re.IGNORECASE)
|
|
|
|
xmax_tag = pattern.findall(
|
|
|
|
str(ET.tostringlist(box_element)))[0][1:-1]
|
|
|
|
x2 = float(box_element.find(xmax_tag).text)
|
|
|
|
pattern = re.compile('<ymax>', re.IGNORECASE)
|
|
|
|
ymax_tag = pattern.findall(
|
|
|
|
str(ET.tostringlist(box_element)))[0][1:-1]
|
|
|
|
y2 = float(box_element.find(ymax_tag).text)
|
|
|
|
x1 = max(0, x1)
|
|
|
|
y1 = max(0, y1)
|
|
|
|
if im_w > 0.5 and im_h > 0.5:
|
|
|
|
x2 = min(im_w - 1, x2)
|
|
|
|
y2 = min(im_h - 1, y2)
|
|
|
|
|
|
|
|
if not (x2 >= x1 and y2 >= y1):
|
|
|
|
logging.warning(
|
|
|
|
"Bounding box for object {} does not satisfy xmin {} <= xmax {} and ymin {} <= ymax {}, "
|
|
|
|
"so this object is skipped. xml file: {}".format(
|
|
|
|
i, x1, x2, y1, y2, xml_file))
|
|
|
|
continue
|
|
|
|
|
|
|
|
gt_bbox[i, :] = [x1, y1, x2, y2]
|
|
|
|
gt_class[i, 0] = cname2cid[cname]
|
|
|
|
gt_score[i, 0] = 1.
|
|
|
|
is_crowd[i, 0] = 0
|
|
|
|
difficult[i, 0] = _difficult
|
|
|
|
i += 1
|
|
|
|
annotations['annotations'].append({
|
|
|
|
'iscrowd': 0,
|
|
|
|
'image_id': int(im_id[0]),
|
|
|
|
'bbox': [x1, y1, x2 - x1, y2 - y1],
|
|
|
|
'area': float((x2 - x1) * (y2 - y1)),
|
|
|
|
'category_id': cname2cid[cname] + 1,
|
|
|
|
'id': ann_ct,
|
|
|
|
'difficult': _difficult
|
|
|
|
})
|
|
|
|
ann_ct += 1
|
|
|
|
|
|
|
|
gt_bbox = gt_bbox[:i, :]
|
|
|
|
gt_class = gt_class[:i, :]
|
|
|
|
gt_score = gt_score[:i, :]
|
|
|
|
is_crowd = is_crowd[:i, :]
|
|
|
|
difficult = difficult[:i, :]
|
|
|
|
|
|
|
|
im_info = {
|
|
|
|
'im_id': im_id,
|
|
|
|
'image_shape': np.array(
|
|
|
|
[im_h, im_w], dtype=np.int32)
|
|
|
|
}
|
|
|
|
label_info = {
|
|
|
|
'is_crowd': is_crowd,
|
|
|
|
'gt_class': gt_class,
|
|
|
|
'gt_bbox': gt_bbox,
|
|
|
|
'gt_score': gt_score,
|
|
|
|
'difficult': difficult
|
|
|
|
}
|
|
|
|
|
|
|
|
if gt_bbox.size > 0:
|
|
|
|
self.file_list.append({
|
|
|
|
'image': img_file,
|
|
|
|
**
|
|
|
|
im_info,
|
|
|
|
**
|
|
|
|
label_info
|
|
|
|
})
|
|
|
|
annotations['images'].append({
|
|
|
|
'height': im_h,
|
|
|
|
'width': im_w,
|
|
|
|
'id': int(im_id[0]),
|
|
|
|
'file_name': osp.split(img_file)[1]
|
|
|
|
})
|
|
|
|
else:
|
|
|
|
neg_file_list.append({
|
|
|
|
'image': img_file,
|
|
|
|
**
|
|
|
|
im_info,
|
|
|
|
**
|
|
|
|
label_info
|
|
|
|
})
|
|
|
|
ct += 1
|
|
|
|
|
|
|
|
if self.use_mix:
|
|
|
|
self.num_max_boxes = max(self.num_max_boxes, 2 * len(objs))
|
|
|
|
else:
|
|
|
|
self.num_max_boxes = max(self.num_max_boxes, len(objs))
|
|
|
|
|
|
|
|
if not ct:
|
|
|
|
logging.error("No voc record found in %s' % (file_list)", exit=True)
|
|
|
|
self.pos_num = len(self.file_list)
|
|
|
|
if self.allow_empty and neg_file_list:
|
|
|
|
self.file_list += self._sample_empty(neg_file_list)
|
|
|
|
logging.info(
|
|
|
|
"{} samples in file {}, including {} positive samples and {} negative samples.".
|
|
|
|
format(
|
|
|
|
len(self.file_list), file_list, self.pos_num,
|
|
|
|
len(self.file_list) - self.pos_num))
|
|
|
|
self.num_samples = len(self.file_list)
|
|
|
|
self.coco_gt = COCO()
|
|
|
|
self.coco_gt.dataset = annotations
|
|
|
|
self.coco_gt.createIndex()
|
|
|
|
|
|
|
|
self._epoch = 0
|
|
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
|
|
sample = copy.deepcopy(self.file_list[idx])
|
|
|
|
if self.data_fields is not None:
|
|
|
|
sample = {k: sample[k] for k in self.data_fields}
|
|
|
|
if self.use_mix and (self.mixup_op.mixup_epoch == -1 or
|
|
|
|
self._epoch < self.mixup_op.mixup_epoch):
|
|
|
|
if self.num_samples > 1:
|
|
|
|
mix_idx = random.randint(1, self.num_samples - 1)
|
|
|
|
mix_pos = (mix_idx + idx) % self.num_samples
|
|
|
|
else:
|
|
|
|
mix_pos = 0
|
|
|
|
sample_mix = copy.deepcopy(self.file_list[mix_pos])
|
|
|
|
if self.data_fields is not None:
|
|
|
|
sample_mix = {k: sample_mix[k] for k in self.data_fields}
|
|
|
|
sample = self.mixup_op(sample=[
|
|
|
|
ImgDecoder(to_rgb=False)(sample),
|
|
|
|
ImgDecoder(to_rgb=False)(sample_mix)
|
|
|
|
])
|
|
|
|
sample = self.transforms(sample)
|
|
|
|
return sample
|
|
|
|
|
|
|
|
def __len__(self):
|
|
|
|
return self.num_samples
|
|
|
|
|
|
|
|
def set_epoch(self, epoch_id):
|
|
|
|
self._epoch = epoch_id
|
|
|
|
|
|
|
|
def cluster_yolo_anchor(self,
|
|
|
|
num_anchors,
|
|
|
|
image_size,
|
|
|
|
cache=True,
|
|
|
|
cache_path=None,
|
|
|
|
iters=300,
|
|
|
|
gen_iters=1000,
|
|
|
|
thresh=.25):
|
|
|
|
"""
|
|
|
|
Cluster YOLO anchors.
|
|
|
|
|
|
|
|
Reference:
|
|
|
|
https://github.com/ultralytics/yolov5/blob/master/utils/autoanchor.py
|
|
|
|
|
|
|
|
Args:
|
|
|
|
num_anchors (int): number of clusters
|
|
|
|
image_size (list or int): [h, w], being an int means image height and image width are the same.
|
|
|
|
cache (bool): whether using cache
|
|
|
|
cache_path (str or None, optional): cache directory path. If None, use `data_dir` of dataset.
|
|
|
|
iters (int, optional): iters of kmeans algorithm
|
|
|
|
gen_iters (int, optional): iters of genetic algorithm
|
|
|
|
threshold (float, optional): anchor scale threshold
|
|
|
|
verbose (bool, optional): whether print results
|
|
|
|
"""
|
|
|
|
if cache_path is None:
|
|
|
|
cache_path = self.data_dir
|
|
|
|
cluster = YOLOAnchorCluster(
|
|
|
|
num_anchors=num_anchors,
|
|
|
|
dataset=self,
|
|
|
|
image_size=image_size,
|
|
|
|
cache=cache,
|
|
|
|
cache_path=cache_path,
|
|
|
|
iters=iters,
|
|
|
|
gen_iters=gen_iters,
|
|
|
|
thresh=thresh)
|
|
|
|
anchors = cluster()
|
|
|
|
return anchors
|
|
|
|
|
|
|
|
def add_negative_samples(self, image_dir, empty_ratio=1):
|
|
|
|
"""将背景图片加入训练
|
|
|
|
|
|
|
|
Args:
|
|
|
|
image_dir (str):背景图片所在的文件夹目录。
|
|
|
|
empty_ratio (float or None): 用于指定负样本占总样本数的比例。如果为None,保留数据集初始化是设置的`empty_ratio`值,
|
|
|
|
否则更新原有`empty_ratio`值。如果小于0或大于等于1,则保留全部的负样本。默认为1。
|
|
|
|
|
|
|
|
"""
|
|
|
|
import cv2
|
|
|
|
if not osp.isdir(image_dir):
|
|
|
|
raise Exception("{} is not a valid image directory.".format(
|
|
|
|
image_dir))
|
|
|
|
if empty_ratio is not None:
|
|
|
|
self.empty_ratio = empty_ratio
|
|
|
|
image_list = os.listdir(image_dir)
|
|
|
|
max_img_id = max(len(self.file_list) - 1, max(self.coco_gt.getImgIds()))
|
|
|
|
neg_file_list = list()
|
|
|
|
for image in image_list:
|
|
|
|
if not is_pic(image):
|
|
|
|
continue
|
|
|
|
gt_bbox = np.zeros((0, 4), dtype=np.float32)
|
|
|
|
gt_class = np.zeros((0, 1), dtype=np.int32)
|
|
|
|
gt_score = np.zeros((0, 1), dtype=np.float32)
|
|
|
|
is_crowd = np.zeros((0, 1), dtype=np.int32)
|
|
|
|
difficult = np.zeros((0, 1), dtype=np.int32)
|
|
|
|
|
|
|
|
max_img_id += 1
|
|
|
|
im_fname = osp.join(image_dir, image)
|
|
|
|
img_data = cv2.imread(im_fname, cv2.IMREAD_UNCHANGED)
|
|
|
|
im_h, im_w, im_c = img_data.shape
|
|
|
|
|
|
|
|
im_info = {
|
|
|
|
'im_id': np.asarray([max_img_id]),
|
|
|
|
'image_shape': np.array(
|
|
|
|
[im_h, im_w], dtype=np.int32)
|
|
|
|
}
|
|
|
|
label_info = {
|
|
|
|
'is_crowd': is_crowd,
|
|
|
|
'gt_class': gt_class,
|
|
|
|
'gt_bbox': gt_bbox,
|
|
|
|
'gt_score': gt_score,
|
|
|
|
'difficult': difficult
|
|
|
|
}
|
|
|
|
if 'gt_poly' in self.file_list[0]:
|
|
|
|
label_info['gt_poly'] = []
|
|
|
|
|
|
|
|
neg_file_list.append({'image': im_fname, ** im_info, ** label_info})
|
|
|
|
if neg_file_list:
|
|
|
|
self.allow_empty = True
|
|
|
|
self.file_list += self._sample_empty(neg_file_list)
|
|
|
|
logging.info(
|
|
|
|
"{} negative samples added. Dataset contains {} positive samples and {} negative samples.".
|
|
|
|
format(
|
|
|
|
len(self.file_list) - self.num_samples, self.pos_num,
|
|
|
|
len(self.file_list) - self.pos_num))
|
|
|
|
self.num_samples = len(self.file_list)
|
|
|
|
|
|
|
|
def _sample_empty(self, neg_file_list):
|
|
|
|
if 0. <= self.empty_ratio < 1.:
|
|
|
|
import random
|
|
|
|
total_num = len(self.file_list)
|
|
|
|
neg_num = total_num - self.pos_num
|
|
|
|
sample_num = min((total_num * self.empty_ratio - neg_num) //
|
|
|
|
(1 - self.empty_ratio), len(neg_file_list))
|
|
|
|
return random.sample(neg_file_list, sample_num)
|
|
|
|
else:
|
|
|
|
return neg_file_list
|