import numpy as np
from ndsampler import abstract_sampler
from ndsampler import category_tree
import networkx as nx
import kwarray
import kwimage
# Moved to kwcoco
from kwcoco.demo.toypatterns import CategoryPatterns
from kwcoco.demo.toydata import demodata_toy_img, demodata_toy_dset # NOQA
[docs]
class DynamicToySampler(abstract_sampler.AbstractSampler):
"""
Generates positive and negative samples on the fly.
Note:
Its probably more robust to generate a static fixed-size dataset with
'demodata_toy_dset' or `kwcoco.CocoDataset.demo`. However, if you
need a sampler that dynamically generates toydata, this is for you.
Ignore:
>>> from ndsampler.toydata import *
>>> self = DynamicToySampler()
>>> window_dims = (96, 96)
img, anns = self.load_positive(window_dims=window_dims)
kwplot.autompl()
kwplot.imshow(img['imdata'])
img, anns = self.load_negative(window_dims=window_dims)
kwplot.autompl()
kwplot.imshow(img['imdata'])
CommandLine:
xdoctest -m ndsampler.toydata DynamicToySampler --show
Example:
>>> # Test that this sampler works with the dataset
>>> from ndsampler.toydata import *
>>> self = DynamicToySampler(1e3)
>>> imgs = [self.load_positive()['im'] for _ in range(9)]
>>> # xdoctest: +REQUIRES(--show)
>>> stacked = kwimage.stack_images_grid(imgs, overlap=-10)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(stacked)
>>> kwplot.show_if_requested()
"""
def __init__(self, n_positives=1e5, seed=None, gsize=(416, 416),
categories=None):
self.catpats = CategoryPatterns.coerce(categories)
self.kp_catnames = self.catpats.kp_catnames
self.cname_to_cid = {
cat['name']: cat['id'] for cat in self.catpats
}
self.cname_to_cid['background'] = 0
self.cid_to_cname = {
cid: cname for cname, cid in self.cname_to_cid.items()
}
graph = nx.DiGraph()
graph.add_nodes_from(self.cname_to_cid.keys())
nx.set_node_attributes(graph, self.cname_to_cid, name='id')
self.catgraph = category_tree.CategoryTree(graph)
self.BACKGROUND_CLASS_ID = self.catgraph.node_to_id['background']
self._full_imgsize = gsize
self._n_positives = int(n_positives)
self._n_images = 50
self.seed = seed
self.gray = True
self._n_annots_pos = (1, 100)
self._n_annots_neg = (0, 10)
# self.catpats = CategoryPatterns.coerce(self.catpats)
# fg_scale=fg_scale, fg_intensity=fg_intensity, rng=rng)
[docs]
def load_item(self, index, pad=None, window_dims=None):
"""
Loads from positives and then negatives.
"""
if index < self._n_positives:
sample = self.load_positive(index, pad=pad, window_dims=window_dims)
else:
index = index - self._n_positives
sample = self.load_negative(index, pad=pad, window_dims=window_dims)
return sample
def __len__(self):
return self._n_positives * 2
def _depends(self):
return [self._n_positives]
@property
def class_ids(self):
return list(self.cid_to_cname.keys())
@property
def n_positives(self):
return self._n_positives
@property
def n_annots(self):
raise NotImplementedError
@property
def n_images(self):
raise NotImplementedError
[docs]
def image_ids(self):
# All images are dummy images
return list(range(10))
[docs]
def lookup_class_name(self, class_id):
return self.cid_to_cname[class_id]
[docs]
def lookup_class_id(self, class_name):
return self.cname_to_cid[class_name]
def _lookup_kpnames(self, class_id):
cname = self.lookup_class_name(class_id)
return self.catpats.cname_to_kp[cname]
# if cname is not None:
# return ['center']
@property
def n_categories(self):
return len(self.catpats) + 1
[docs]
def preselect(self, n_pos=None, n_neg=None, neg_to_pos_ratio=None,
window_dims=None, rng=None, verbose=0):
n_pos = self._n_positives
if n_neg is None:
n_neg = int(n_pos * neg_to_pos_ratio)
if verbose:
print('no need to presample dynamic sampler')
return n_pos, n_neg
[docs]
def load_image(self, image_id=None, rng=None):
img, anns = self.load_image_with_annots(image_id=image_id, rng=rng)
return img['imdata']
[docs]
def load_image_with_annots(self, image_id=None, rng=None):
""" Returns a random image and its annotations """
if image_id is not None and self.seed is not None:
rng = kwarray.ensure_rng(self.seed * len(self) + image_id)
rng = kwarray.ensure_rng(rng)
img, anns = demodata_toy_img(gsize=self._full_imgsize,
categories=self.catpats,
gray=self.gray,
rng=rng, n_annots=(0, 10))
_node_to_id = self.catgraph.node_to_id
img['id'] = int(rng.rand() * self._n_images)
img['file_name'] = '{}.png'.format(img['id'])
for ann in anns:
ann['category_id'] = _node_to_id[ann['category_name']]
return img, anns
[docs]
def load_sample(self, tr, pad=None, window_dims=None):
raise NotImplementedError
def _load_toy_sample(self, window_dims, pad, rng, centerobj, n_annots):
rng = kwarray.ensure_rng(rng)
gid = int(rng.rand() * 500)
# guuid = ub.hash_data(rng)
if window_dims is None:
window_dims = self._full_imgsize[::-1]
gsize = np.array(window_dims[::-1])
if pad is not None:
gsize += 2 * np.array(pad[::-1])
img, anns = demodata_toy_img(gsize=gsize,
rng=rng,
categories=self.catpats,
gray=self.gray,
centerobj=centerobj, n_annots=n_annots)
im = img['imdata']
if centerobj == 'neg':
cid = self.BACKGROUND_CLASS_ID
aid = -1
else:
ann = anns[0]
cname = ann['category_name']
cid = self.cname_to_cid[cname]
aid = 1
tr_ = {
'aid': aid,
# 'gid': guuid,
# 'id': gid,
'gid': gid,
'rel_cx': gsize[0] / 2,
'rel_cy': gsize[1] / 2,
'cx': gsize[0] / 2,
'cy': gsize[1] / 2,
'width': window_dims[1],
'height': window_dims[0],
'category_id': cid,
}
# Handle segmentations and keypoints if they exist
sseg_list = []
kpts_list = []
kp_catnames = self.catpats.kp_catnames
for ann in anns:
coco_sseg = ann.get('segmentation', None)
coco_kpts = ann.get('keypoints', None)
if coco_kpts is not None:
cid = self.lookup_class_id(ann['category_name'])
# kpnames = self._lookup_kpnames(cid)
rel_points = kwimage.Points._from_coco(coco_kpts,
classes=kp_catnames)
# coco_xyf = np.array(coco_kpts).reshape(-1, 3)
# flags = (coco_xyf.T[2] > 0)
# xy_pts = coco_xyf[flags, 0:2]
# kpnames = list(ub.compress(kpnames, flags))
# kp_class_idxs = np.array([kp_catnames.index(n) for n in kpnames])
# rel_points = kwimage.Points(xy=xy_pts,
# class_idxs=kp_class_idxs,
# classes=kp_catnames)
else:
rel_points = None
if coco_sseg is not None:
# TODO: implement MultiPolygon coerce instead
data_dims = gsize[::-1]
abs_sseg = kwimage.Mask.coerce(coco_sseg, dims=data_dims)
rel_sseg = abs_sseg.to_multi_polygon()
else:
rel_sseg = None
kpts_list.append(rel_points)
sseg_list.append(rel_sseg)
rel_ssegs = kwimage.PolygonList(sseg_list)
rel_kpts = kwimage.PolygonList(kpts_list)
rel_kpts.meta['classes'] = self.catpats.kp_catnames
rel_boxes = kwimage.Boxes([a['bbox'] for a in anns], 'xywh').to_cxywh()
annots = {
'aids': np.arange(len(anns)),
'cids': np.array([self.lookup_class_id(a['category_name']) for a in anns]),
'rel_cxywh': rel_boxes.data,
'rel_boxes': rel_boxes,
'rel_kpts': rel_kpts,
'rel_ssegs': rel_ssegs,
}
# TODO: if window_dims was None, then crop the image to the size of the
# annotation and remove non-visible other annots!
sample = {'im': im, 'target': tr_, 'annots': annots}
return sample
[docs]
def load_positive(self, index=None, pad=None, window_dims=None, rng=None):
"""
Note: window_dims is height / width
Example:
>>> from ndsampler.toydata import *
>>> self = DynamicToySampler(1e2)
>>> sample = self.load_positive()
>>> annots = sample['annots']
>>> assert len(annots['aids']) > 0
>>> assert len(annots['rel_cxywh']) == len(annots['aids'])
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> # Draw box in relative sample context
>>> kwplot.imshow(sample['im'], pnum=(1, 1, 1), fnum=1)
>>> annots['rel_boxes'].translate([-.5, -.5]).draw()
>>> annots['rel_ssegs'].draw(color='red', alpha=.6)
>>> annots['rel_kpts'].draw(color='green', alpha=.8, radius=4)
"""
if index is not None and self.seed is not None:
rng = self.seed * len(self) + index
sample = self._load_toy_sample(window_dims, pad, rng,
centerobj='pos',
n_annots=self._n_annots_pos)
return sample
[docs]
def load_negative(self, index=None, pad=None, window_dims=None, rng=None):
if index is not None and self.seed is not None:
rng = kwarray.ensure_rng(self.seed * len(self) + index)
sample = self._load_toy_sample(window_dims, pad, rng,
centerobj='neg',
n_annots=self._n_annots_neg)
return sample