You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

442 lines
18 KiB

# 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 .base import BaseDataset
from paddlers.utils import logging, get_encoding, norm_path, is_pic
from paddlers.transforms import DecodeImg, MixupImage
from paddlers.tools import YOLOAnchorCluster
class VOCDetDataset(BaseDataset):
"""
Dataset with PASCAL VOC annotations for detection tasks.
Args:
data_dir (str): Root directory of the dataset.
file_list (str): Path of the file that contains relative paths of images and annotation files.
transforms (paddlers.transforms.Compose): Data preprocessing and data augmentation operators to apply.
label_list (str|None, optional): Path of the file that contains the category names. Defaults to None.
num_workers (int|str, optional): Number of processes used for data loading. If `num_workers` is 'auto',
the number of workers will be automatically determined according to the number of CPU cores: If
there are more than 16 cores,8 workers will be used. Otherwise, the number of workers will be half
the number of CPU cores. Defaults: 'auto'.
shuffle (bool, optional): Whether to shuffle the samples. Defaults to False.
allow_empty (bool, optional): Whether to add negative samples. Defaults to False.
empty_ratio (float, optional): Ratio of negative samples. If `empty_ratio` is smaller than 0 or not less
than 1, keep all generated negative samples. Defaults to 1.0.
"""
def __init__(self,
data_dir,
file_list,
transforms,
label_list,
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.
import matplotlib
matplotlib.use('Agg')
from pycocotools.coco import COCO
super(VOCDetDataset, self).__init__(data_dir, label_list, transforms,
num_workers, shuffle)
self.data_fields = None
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.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 ValueError("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 = norm_path(img_file)
xml_file = norm_path(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 is no field '<bndbox>' in the 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=[
DecodeImg(to_rgb=False)(sample),
DecodeImg(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[int]|int): [h, w] or an int value that corresponds to the shape [image_size, image_size].
cache (bool, optional): Whether to use cache. Defaults to True.
cache_path (str|None, optional): Path of cache directory. If None, use `dataset.data_dir`.
Defaults to None.
iters (int, optional): Iterations of k-means algorithm. Defaults to 300.
gen_iters (int, optional): Iterations of genetic algorithm. Defaults to 1000.
thresh (float, optional): Anchor scale threshold. Defaults to 0.25.
"""
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):
"""
Generate and add negative samples.
Args:
image_dir (str): Directory that contains images.
empty_ratio (float|None, optional): Ratio of negative samples. If `empty_ratio` is smaller than
0 or not less than 1, keep all generated negative samples. Defaults to 1.0.
"""
import cv2
if not osp.isdir(image_dir):
raise ValueError("{} 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