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}