text: import time
  import json
  import datetime
  import os
  import tensorflow.compat.v1 as tf
  import math
  import glob
  import numpy as np
  import itertools
  import setproctitle
  import argparse
  from pathlib import Path
  tf.enable_eager_execution()
  
  from waymo_open_dataset.utils import range_image_utils
  from waymo_open_dataset.utils import transform_utils
  from waymo_open_dataset.utils import  frame_utils
  from waymo_open_dataset import dataset_pb2 as open_dataset
  
  from pycocotools import mask
  from mask_utils import mask2polygon, segm2json, xyxy2xywh
  from mmdet.apis import inference_detector, init_detector
  
  INFO = {
      "description": "Waymo Masks",
      "url": "",
      "version": "0.1.0",
      "year": 2021,
      "contributor": "darrenjkt",
      "date_created": datetime.datetime.utcnow().isoformat(' ')
  }
  
  LICENSES = [
      {
          "id": 1,
          "name": "Attribution-NonCommercial-ShareAlike License",
          "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/"
      }
  ]
  
  CATEGORIES = [
      {
          'id': 0,
          'name': 'person',
          'supercategory': 'person',
      },
      {
          'id': 1,
          'name': 'bicycle',
          'supercategory': 'vehicle',
      },
      {
          'id': 2,
          'name': 'car',
          'supercategory': 'vehicle',
      },
  ]
  
  coco_output = {
          "info": INFO,
          "licenses": LICENSES,
          "categories": CATEGORIES,
          "images": [],
          "annotations": []
      }
  
  
  def create_image_info(image_id, file_name, image_size, 
                        date_captured=datetime.datetime.utcnow().isoformat(' ')):
  
      image_info = {
              "id": image_id,
              "file_name": file_name,
              "width": image_size[1],
              "height": image_size[0],
              "date_captured": date_captured,
      }
  
      return image_info
  
  def parse_args():
      parser = argparse.ArgumentParser(
          description='Get segmentation masks for image dataset and save to COCO json annotation format')
      parser.add_argument('--config', required=True, help='Model config file path')
      parser.add_argument('--checkpoint', required=True, help='Trained model checkpoint file')
      parser.add_argument('--camera', required=True, help='Specify which camera to use from the list [FRONT,FRONT_LEFT,FRONT_RIGHT,SIDE_LEFT,SIDE_RIGHT]')
      parser.add_argument('--data_dir', required=True, help='Root folder of waymo data')
      parser.add_argument('--output_json', required=True, help='Absolute path and filename to save json annotation files')
      args = parser.parse_args()
  
      print("\n----- Args -----")
      for arg in vars(args):
          print (f'- {arg}: {getattr(args, arg)}')
      print("\n")
  
      return args
  
      
  if __name__ == "__main__":
  
      start_time = time.time()
      args = parse_args()    
      data_dir = Path(args.data_dir)
      
      camera_names = ['FRONT','FRONT_LEFT','FRONT_RIGHT','SIDE_LEFT','SIDE_RIGHT']
      idx2camera = {v+1:k for v, k in enumerate(camera_names)}
      tf_records = glob.glob(f"{str(data_dir)}/*.tfrecord")
      assert len(tf_records) != 0, print(f"No tfrecords found at: {str(data_dir)}")
  
      model = init_detector(args.config, args.checkpoint, device='cuda:0')
      
      # Process all images in all scenes
      for record_idx in range(0,len(tf_records)):
          start = time.time()
          
          record_fname = tf_records[record_idx]
          record = tf.data.TFRecordDataset(record_fname, compression_type='')
          setproctitle.setproctitle(f'Masks[WAYM]:{record_idx}/{len(tf_records)} records')
          
          for f_idx, data in enumerate(record):
              tic = time.time()
              
              frame = open_dataset.Frame()
              frame.ParseFromString(bytearray(data.numpy()))
              
              camera_images = {idx2camera[v.name]:v for k,v in enumerate(frame.images)}
              image = np.asarray(tf.image.decode_jpeg(camera_images[args.camera].image, dct_method="INTEGER_ACCURATE"))
                 
              # Use the {camera_name}{timestamp in milliseconds converted to int} as token
              image_token = f'{frame.context.name}+{args.camera}+{frame.timestamp_micros}'
              image_info = create_image_info(image_id=image_token, 
                                             file_name=os.path.basename(record_fname), 
                                             image_size=[image.shape[0],image.shape[1]])
              coco_output["images"].append(image_info)
  
              # Get segmentation mask
              result = inference_detector(model, image)
              segm_json_result = segm2json(result, score_thresh=0.3)
              num_predicted = len(segm_json_result)
  
              # Convert to COCO json format    
              for iidx, segm in enumerate(segm_json_result):
  
                  category_info = {'id': segm["category_id"], 'is_crowd': 0}
                  instance_id = f'{image_token}{iidx:0{len(str(num_predicted))}}'
  
                  seg_mask = segm['segmentation']
                  binary_mask_encoded = mask.encode(np.asfortranarray(seg_mask.astype(np.uint8)))
                  area = mask.area(binary_mask_encoded)        
  
                  annotation_info = {
                      "id": instance_id,
                      "image_id": image_token,
                      "category_id": category_info["id"],
                      "iscrowd": category_info["is_crowd"],
                      "area": area.tolist(),
                      "bbox": segm['bbox'], 
                      "segmentation": mask2polygon(seg_mask),
                  } 
  
                  if annotation_info is not None:
                      coco_output["annotations"].append(annotation_info)
  
              print(f'Record: {record_idx+1}/{len(tf_records)}, Image: {f_idx} | Time: {(time.time() - tic):.4}s, Total: {(time.time() - start):.4}s')
  
          print(f"Record {record_idx+1}/{len(tf_records)} Done. \n")
  
      # Save JSON file
      os.makedirs(os.path.dirname(args.output_json),exist_ok=True)
      with open(args.output_json, 'w+') as output_json_file:
          json.dump(coco_output, output_json_file)
  
      print(f'Finished.\nTime taken: {time.time() - start_time}s')
  