| 29 | |
| 30 | |
| 31 | class TFRecordDetectionPipeline(Pipeline): |
| 32 | def __init__(self, args): |
| 33 | super(TFRecordDetectionPipeline, self).__init__(args.batch_size, args.num_workers, 0, 0) |
| 34 | self.input = ops.readers.TFRecord( |
| 35 | path=os.path.join(test_dummy_data_path, "small_coco.tfrecord"), |
| 36 | index_path=os.path.join(test_dummy_data_path, "small_coco_index.idx"), |
| 37 | features={ |
| 38 | "image/encoded": tfrec.FixedLenFeature((), tfrec.string, ""), |
| 39 | "image/object/class/label": tfrec.VarLenFeature([], tfrec.int64, 0), |
| 40 | "image/object/bbox": tfrec.VarLenFeature([4], tfrec.float32, 0.0), |
| 41 | }, |
| 42 | shard_id=0, |
| 43 | num_shards=1, |
| 44 | random_shuffle=False, |
| 45 | ) |
| 46 | |
| 47 | self.decode_gpu = ops.decoders.Image(device="mixed", output_type=types.RGB) |
| 48 | self.cast = ops.Cast(dtype=types.INT32) |
| 49 | self.box_encoder = ops.BoxEncoder(device="cpu", criteria=0.5, anchors=coco_anchors()) |
| 50 | |
| 51 | def define_graph(self): |
| 52 | inputs = self.input() |
| 53 | input_images = inputs["image/encoded"] |
| 54 | |
| 55 | image_gpu = self.decode_gpu(input_images) |
| 56 | labels = self.cast(inputs["image/object/class/label"]) |
| 57 | encoded_boxes, encoded_labels = self.box_encoder(inputs["image/object/bbox"], labels) |
| 58 | |
| 59 | return (image_gpu, inputs["image/object/bbox"], labels, encoded_boxes, encoded_labels) |
| 60 | |
| 61 | |
| 62 | class COCODetectionPipeline(Pipeline): |