def __init__(self,
root,
annFile,
image_transform=None,
target_transforms=None,
n_images=None,
preprocess=None,
all_images=False,
all_persons=False):
from pycocotools.coco import COCO
self.root = root
self.coco = COCO(annFile)
# got category names and corresponding cat_ids
catNms = ['car']
self.cat_ids = self.coco.getCatIds(catNms=catNms)
#print('category length'+str(self.cat_ids))
#reverse cocolabels to (name:id)
self.reverse_dict = {value: key for key, value in COCO_LABELS.items()}
# new_cat_id is the id corresponding to the choosen category name
self.new_cat_id = self.reverse_dict[catNms[0]]
#print('new category length'+str(self.new_cat_id))
# self.dict is used to map original cat_id to our cat_id
#self.dict = dict(zip(self.cat_ids, self.new_cat_id))
self.dict = {self.cat_ids[0]: self.new_cat_id}
#print(type(self.dict))
#self.ids only contains the images_id of which image contain the choosen category object
#self.ids = []
#self.ids.append(self.coco.getImgIds(catIds = self.cat_ids))
self.ids = self.coco.getImgIds(catIds=self.cat_ids)
# set of all image_ids of coco dataset
self.all_ids_set = set(self.coco.getImgIds())
# set of all choosen image ids
self.ids_set = set(self.ids)
# image_ids which does not contain the object which I choose
self.other_obejct_id = list(self.all_ids_set - self.ids_set)
# choose some images which do not contain the object
num_images = int(0.8 * len(self.ids))
self.other_id = sample(self.other_obejct_id, num_images)
#merge two lists of ids
self.ids = self.ids + self.other_id
if n_images:
self.ids = self.ids[:n_images]
print('Images: {}'.format(len(self.ids)))
self.preprocess = preprocess or transforms.Normalize()
self.image_transform = image_transform or transforms.image_transform
self.target_transforms = target_transforms
self.log = logging.getLogger(self.__class__.__name__)
def rescale_from_12_to_6(im_np, x, y):
print(im_np[0, :, 0])
im = PIL.Image.fromarray(im_np)
anns = [{
'keypoints': [[x, y, 2.0]],
'bbox': [0.0, 0.0, 12.0, 12.0],
'segmentation': None,
}]
transform = transforms.Compose([
transforms.Normalize(),
transforms.RescaleAbsolute(6),
])
return transform(im, anns, None)
def __init__(self, root, annFile, image_transform=None, target_transforms=None,
n_images=None, preprocess=None, all_images=False, all_persons=False):
from pycocotools.coco import COCO
self.root = root
self.coco = COCO(annFile)
#print(self.root)
# got all 80 category names and corresponding cat_ids
catNms = list(COCO_LABELS.values())[1:]
self.cat_ids = self.coco.getCatIds(catNms=catNms)
# self.dict is used to map original cat_id to our cat_id
new_cat_id = [i+1 for i in range(len(catNms))]
self.dict = dict(zip(self.cat_ids, new_cat_id))
self.ids = self.coco.getImgIds()
#self.ids = []
#for i in range(len(self.cat_ids)):
#self.ids.append(self.coco.getImgIds(catIds = self.cat_ids[i]))
# if all_images:
# self.ids = self.coco.getImgIds(catIds =self.cat_ids)
# #print(self.ids)
# elif all_persons:
# self.ids = self.coco.getImgIds(catIds=self.cat_ids)
# else:
# self.ids = self.coco.getImgIds(catIds=self.cat_ids)
# #self.filter_for_keypoint_annotations()
if n_images:
self.ids = self.ids[:n_images]
print('Images: {}'.format(len(self.ids)))
self.preprocess = preprocess or transforms.Normalize()
self.image_transform = image_transform or transforms.image_transform
self.target_transforms = target_transforms
self.log = logging.getLogger(self.__class__.__name__)
def generate(m, inputs):
args = cli()
model, processor = m
image = inputs["image"]
# data
preprocess = None
if args.long_edge:
preprocess = transforms.Compose([
transforms.Normalize(),
transforms.RescaleAbsolute(args.long_edge),
transforms.CenterPad(args.long_edge),
])
data = datasets.PilImageList([image], preprocess=preprocess)
data_loader = torch.utils.data.DataLoader(data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=args.pin_memory,
num_workers=args.loader_workers)
# visualizers
keypoint_painter = show.KeypointPainter(show_box=False)
skeleton_painter = show.KeypointPainter(show_box=False,
color_connections=True,
markersize=1,
linewidth=6)
image_paths, image_tensors, processed_images_cpu = next(iter(data_loader))
images = image_tensors.permute(0, 2, 3, 1)
processed_images = processed_images_cpu.to(args.device, non_blocking=True)
fields_batch = processor.fields(processed_images)
pred_batch = processor.annotations_batch(fields_batch,
debug_images=processed_images_cpu)
# unbatch
image_path, image, processed_image_cpu, pred = image_paths[0], images[
0], processed_images_cpu[0], pred_batch[0]
processor.set_cpu_image(image, processed_image_cpu)
keypoint_sets, scores = processor.keypoint_sets_from_annotations(pred)
kp_json = json.dumps([{
'keypoints': np.around(kps, 1).reshape(-1).tolist(),
'bbox': bbox_from_keypoints(kps),
} for kps in keypoint_sets])
kwargs = {
'figsize': (args.figure_width,
args.figure_width * image.shape[0] / image.shape[1]),
}
fig = plt.figure(**kwargs)
ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])
ax.set_axis_off()
ax.set_xlim(0, image.shape[1])
ax.set_ylim(image.shape[0], 0)
fig.add_axes(ax)
ax.imshow(image)
skeleton_painter.keypoints(ax, keypoint_sets, scores=scores)
fig.canvas.draw()
w, h = fig.canvas.get_width_height()
output_image = np.fromstring(fig.canvas.tostring_rgb(),
dtype='uint8').reshape(h, w, 3)
return {'keypoints': kp_json, 'image': output_image}