| 60 | |
| 61 | |
| 62 | class COCODetectionPipeline(Pipeline): |
| 63 | def __init__(self, args, data_path=test_data_path): |
| 64 | super(COCODetectionPipeline, self).__init__(args.batch_size, args.num_workers, 0, 0) |
| 65 | |
| 66 | self.input = ops.readers.COCO( |
| 67 | file_root=os.path.join(data_path, "images"), |
| 68 | annotations_file=os.path.join(data_path, "instances.json"), |
| 69 | shard_id=0, |
| 70 | num_shards=1, |
| 71 | ratio=True, |
| 72 | ltrb=True, |
| 73 | random_shuffle=False, |
| 74 | ) |
| 75 | |
| 76 | self.decode_gpu = ops.decoders.Image(device="mixed", output_type=types.RGB) |
| 77 | self.box_encoder = ops.BoxEncoder(device="cpu", criteria=0.5, anchors=coco_anchors()) |
| 78 | |
| 79 | def define_graph(self): |
| 80 | inputs, boxes, labels = self.input(name="Reader") |
| 81 | image_gpu = self.decode_gpu(inputs) |
| 82 | encoded_boxes, encoded_labels = self.box_encoder(boxes, labels) |
| 83 | |
| 84 | return (image_gpu, boxes, labels, encoded_boxes, encoded_labels) |
| 85 | |
| 86 | |
| 87 | def print_args(args): |