class CocoDetection(Dataset):
def __init__(self, root, ann_file, transform=None):
from hpycocotools.coco import COCO
self.root = root
self.ann_file = ann_file
with open(self.ann_file, 'r') as f:
self.data = json.load(f)
self.coco = COCO(self.data, verbose=False)
self.ids = list(self.coco.imgs.keys())
self.transform = transform
def to_coco(self, indices=None):
if indices is None:
return self.data
ids = [self.ids[i] for i in indices]
images = self.coco.loadImgs(ids)
ann_ids = self.coco.getAnnIds(ids)
annotations = self.coco.loadAnns(ann_ids)
return {
**self.data,
"images": images,
"annotations": annotations,
}
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: Tuple (image, target). target is the object returned by ``coco.loadAnns``.
"""
coco = self.coco
img_id = self.ids[index]
ann_ids = coco.getAnnIds(imgIds=img_id)
target = coco.loadAnns(ann_ids)
path = coco.loadImgs(img_id)[0]['file_name']
img = Image.open(os.path.join(self.root, path)).convert('RGB')
if self.transform is not None:
img, target = self.transform(img, target)
return img, target
def __len__(self):
return len(self.ids)
def __repr__(self):
fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
fmt_str += ' Number of datapoints: {}\n'.format(self.__len__())
fmt_str += ' Root Location: {}\n'.format(self.root)
tmp = ' Transforms (if any): '
fmt_str += '{0}{1}\n'.format(
tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
return fmt_str
def find_priors_coco(ds, k=3, max_iter=100, verbose=True):
assert hasattr(ds, "to_coco"), "ds must have `to_coco()` method"
from hpycocotools.coco import COCO
coco = COCO(ds.to_coco(), verbose=False)
bboxes = []
for img_id, anns in coco.imgToAnns.items():
img = coco.loadImgs(img_id)[0]
height = img['height']
width = img['width']
for ann in anns:
l, t, w, h = ann['bbox']
l /= width
t /= height
w /= width
h /= height
bboxes.append([l, t, w, h])
bboxes = np.array(bboxes, dtype=np.float32)
sizes = bboxes[:, 2:]
priors = find_priors_kmeans(sizes, k=k, max_iter=max_iter, verbose=verbose)
priors = torch.stack(sorted(priors, key=lambda x: x[0] * x[1]))
return priors
class COCOEval(Metric):
def __init__(self, annotations, iou_type='bbox', percent_area_ranges=None):
self.annotations = annotations
self.iou_type = iou_type
self.percent_area_ranges = percent_area_ranges
super().__init__()
def reset(self):
from hpycocotools.coco import COCO
self.coco_gt = COCO(self.annotations,
verbose=False,
use_percent=bool(self.percent_area_ranges))
self.res = []
def update(self, output):
target, image_dets, batch_size = get(["target", "preds", "batch_size"],
output)
image_gts = target[0]
for dets, gts in zip(image_dets, image_gts):
image_id = gts[0]['image_id']
for d in dets:
d = {**d, 'image_id': image_id}
self.res.append(d)
def compute(self):
from hpycocotools.coco import COCO
img_ids = list(set([d['image_id'] for d in self.res]))
imgs = self.coco_gt.loadImgs(img_ids)
ann_ids = self.coco_gt.getAnnIds(imgIds=img_ids)
anns = self.coco_gt.loadAnns(ann_ids)
annotations = {
**self.annotations,
'images': imgs,
'annotations': anns,
}
coco_gt = COCO(annotations,
verbose=False,
use_percent=bool(self.percent_area_ranges))
from hpycocotools.cocoeval import COCOeval
from hpycocotools.mask import encode
for dt in self.res:
img = self.coco_gt.imgs[dt['image_id']]
width = img['width']
height = img['height']
l, t, w, h = dt['bbox']
l *= width
t *= height
w *= width
h *= height
dt['bbox'] = [l, t, w, h]
if 'segmentation' in dt and self.iou_type == 'segm':
r = int(l + w)
b = int(t + h)
l = max(0, int(l))
t = max(0, int(t))
r = min(r, width)
b = min(b, height)
w = r - l
h = b - t
m = np.zeros((height, width), dtype=np.uint8)
segm = cv2.resize(dt['segmentation'], (w, h),
interpolation=cv2.INTER_NEAREST)
m[t:b, l:r] = segm
dt['segmentation'] = encode(np.asfortranarray(m))
coco_dt = coco_gt.loadRes(self.res)
ev = COCOeval(coco_gt,
coco_dt,
iouType=self.iou_type,
pAreaRng=self.percent_area_ranges,
verbose=False)
ev.evaluate()
ev.accumulate()
ev.summarize()
return ev.stats[0]
class VOCDetection(Dataset):
"""`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Detection Dataset.
Args:
root (string): Root directory of the VOC Dataset.
year (string, optional): The dataset year, supports years 2007 to 2012.
image_set (string, optional): Select the image_set to use, ``trainval`` or ``test``
download (bool, optional): If true, downloads the dataset from the internet and
puts it in root directory. If dataset is already downloaded, it is not
downloaded again.
(default: alphabetic indexing of VOC's 20 classes).
transform (callable, optional): A function/transform that takes in an PIL image
and returns a transformed version. E.g, ``transforms.RandomCrop``
"""
def __init__(self,
root,
year='2012',
image_set='trainval',
download=False,
transform=None):
self.root = Path(root).expanduser().absolute()
self.year = year
self.image_set = image_set
if image_set == 'test':
dataset_dict = TEST_DATASET_YEAR_DICT
else:
dataset_dict = DATASET_YEAR_DICT
self.url = dataset_dict[year]['url']
self.filename = dataset_dict[year]['filename']
self.md5 = dataset_dict[year]['md5']
base_dir = dataset_dict[year]['base_dir']
self.voc_root = self.root / base_dir
self.image_dir = self.voc_root / 'JPEGImages'
self.ann_file_url = dataset_dict[year]['ann_file_url']
self.ann_file = self.voc_root / ("%s%s.json" % (image_set, year))
if download:
self.download()
from hpycocotools.coco import COCO
with open(self.ann_file, 'r') as f:
self.data = json.load(f)
self.coco = COCO(self.data, verbose=False)
self.ids = list(self.coco.imgs.keys())
self.transform = transform
def to_coco(self, indices=None):
if indices is None:
return self.data
ids = [self.ids[i] for i in indices]
images = self.coco.loadImgs(ids)
ann_ids = self.coco.getAnnIds(ids)
annotations = self.coco.loadAnns(ann_ids)
return {
**self.data,
"images": images,
"annotations": annotations,
}
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: Tuple (image, target). target is the object returned by ``coco.loadAnns``.
"""
coco = self.coco
img_id = self.ids[index]
ann_ids = coco.getAnnIds(imgIds=img_id)
anns = coco.loadAnns(ann_ids)
path = coco.loadImgs([img_id])[0]['file_name']
img = Image.open(self.image_dir / path).convert('RGB')
if self.transform is not None:
img, anns = self.transform(img, anns)
return img, anns
def __len__(self):
return len(self.ids)
def download(self):
import tarfile
if self.voc_root.is_dir() and self.ann_file.exists():
print("Dataset found. Skip download or extract.")
return
if not self.voc_root.is_dir():
download_url(self.url, self.root, self.filename, self.md5)
with tarfile.open(self.root / self.filename, "r") as tar:
tar.extractall(path=self.root)
if not self.ann_file.exists():
google_drive_match = re.match(
r"https://drive.google.com/open\?id=(.*)", self.ann_file_url)
file_id = google_drive_match.group(1)
download_google_drive(file_id, self.voc_root, self.ann_file.name)
def __repr__(self):
fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
fmt_str += ' Year: {}\n'.format(self.year)
fmt_str += ' ImageSet: {}\n'.format(self.image_set)
fmt_str += ' Number of datapoints: {}\n'.format(self.__len__())
fmt_str += ' Root Location: {}\n'.format(self.root)
tmp = ' Transforms (if any): '
fmt_str += '{0}{1}\n'.format(
tmp,
self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
return fmt_str
class SVHNDetection(Dataset):
"""`SVHN <http://ufldl.stanford.edu/housenumbers/>`_ Dataset.
Note: The SVHN dataset assigns the label `10` to the digit `0`. However, in this Dataset,
we assign the label `0` to the digit `0` to be compatible with PyTorch loss functions which
expect the class labels to be in the range `[0, C-1]`
Args:
root (string): Root directory of dataset where directory
``SVHN`` exists.
split (string): One of {'train', 'test', 'extra'}.
Accordingly dataset is selected. 'extra' is Extra training set.
transform (callable, optional): A function/transform that takes in an PIL image
and targets simultaneously and returns a transformed version of them.
download (bool, optional): If true, downloads the dataset from the internet and
puts it in root directory. If dataset is already downloaded, it is not
downloaded again.
"""
def __init__(self, root, split='train', transform=None, download=False):
self.root = Path(root).expanduser().absolute()
self.split = split
self.transform = transform
self.filename = SPLIT_FILES[split]['name']
self.md5 = SPLIT_FILES[split]['md5']
self.url = SPLIT_FILES[split]["url"]
self.img_dir = self.root / split
self.ann_filename = split + ".json"
self.ann_dir = self.root / "annotations"
self.ann_dir.mkdir(exist_ok=True, parents=True)
self.ann_file = self.ann_dir / self.ann_filename
if download:
self.download()
with open(self.ann_file, 'r') as f:
self.data = json.load(f)
from hpycocotools.coco import COCO
self.coco = COCO(self.data, verbose=False)
self.ids = list(self.coco.imgs.keys())
def to_coco(self, indices=None):
if indices is None:
return self.data
ids = [self.ids[i] for i in indices]
images = self.coco.loadImgs(ids)
ann_ids = self.coco.getAnnIds(ids)
annotations = self.coco.loadAnns(ann_ids)
return {
**self.data,
"images": images,
"annotations": annotations,
}
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, anns) where target is a dictionary of the XML tree.
"""
coco = self.coco
img_id = self.ids[index]
ann_ids = coco.getAnnIds(imgIds=img_id)
target = coco.loadAnns(ann_ids)
path = coco.loadImgs(img_id)[0]['file_name']
img = Image.open(self.img_dir / path).convert('RGB')
if self.transform is not None:
img, target = self.transform(img, target)
return img, target
def __len__(self):
return len(self.ids)
def download(self):
if self.img_dir.is_dir() and self.ann_file.exists():
print("Dataset found. Skip download or extract")
return
google_drive_match = re.match(
r"https://drive.google.com/open\?id=(.*)", self.url)
if google_drive_match:
file_id = google_drive_match.group(1)
download_google_drive(file_id, self.root, self.filename, self.md5)
else:
download_url(self.url, self.root, self.filename, self.md5)
file_path = self.root / self.filename
with tarfile.open(file_path, "r") as tar:
tar.extractall(path=self.root)
ann_file = self.root / self.ann_filename
ann_file.rename(self.ann_file)