scripts/infer_segmentation_main.py

   1 #!/usr/bin/env python
   2
   3 import sys
   4 import os
   5 import cv2
   6 import argparse
   7 import datetime
   8 import numpy as np
   9
  10 ################################
  11 from pytorch_jacinto_ai.xnn.utils import str2bool
  12 parser = argparse.ArgumentParser()
  13 parser.add_argument('--save_path', type=str, default=None, help='checkpoint save folder')
  14 parser.add_argument('--gpus', type=int, nargs='*', default=None, help='Base learning rate')
  15 parser.add_argument('--batch_size', type=int, default=None, help='Batch size')
  16 parser.add_argument('--model_name', type=str, default=None, help='model name')
  17 parser.add_argument('--dataset_name', type=str, default=None, help='dataset name')
  18 parser.add_argument('--data_path', type=str, default=None, help='data path')
  19 parser.add_argument('--epoch_size', type=float, default=None, help='epoch size. using a fraction will reduce the data used for one epoch')
  20 parser.add_argument('--img_resize', type=int, nargs=2, default=None, help='img_resize size. for training this will be modified according to rand_scale')
  21 parser.add_argument('--rand_scale', type=float, nargs=2, default=None, help='random scale factors for training')
  22 parser.add_argument('--rand_crop', type=int, nargs=2, default=None, help='random crop for training')
  23 parser.add_argument('--output_size', type=int, nargs=2, default=None, help='output size of the evaluation - prediction/groundtruth. this is not used while training as it blows up memory requirement')
  24 parser.add_argument('--quantize', type=str2bool, default=None, help='Quantize the model')
  25 #parser.add_argument('--model_surgery', type=str, default=None, choices=[None, 'pact2'], help='whether to transform the model after defining')
  26 parser.add_argument('--pretrained', type=str, default=None, help='pretrained model')
  27 parser.add_argument('--bitwidth_weights', type=int, default=None, help='bitwidth for weight quantization')
  28 parser.add_argument('--bitwidth_activations', type=int, default=None, help='bitwidth for activation quantization')
  29 cmds = parser.parse_args()
  30
  31 ################################
  32 # taken care first, since this has to be done before importing pytorch
  33 if 'gpus' in vars(cmds):
  34     value = getattr(cmds, 'gpus')
  35     if (value is not None) and ("CUDA_VISIBLE_DEVICES" not in os.environ):
  36         os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(v) for v in value])
  37     #
  38 #
  39
  40 # to avoid hangs in data loader with multi threads
  41 # this was observed after using cv2 image processing functions
  42 # https://github.com/pytorch/pytorch/issues/1355
  43 cv2.setNumThreads(0)
  44
  45
  46 ################################
  47 #import of torch should be after CUDA_VISIBLE_DEVICES for it to take effect
  48 import torch
  49 from pytorch_jacinto_ai.engine import infer_pixel2pixel
  50
  51 # Create the parse and set default arguments
  52 args = infer_pixel2pixel.get_config()
  53
  54
  55 #Modify arguments
  56 args.model_name = "deeplabv3lite_mobilenetv2_tv" #"deeplabv3lite_mobilenetv2_relu" #"deeplabv3lite_mobilenetv2_relu_x1p5" #"deeplabv3plus"
  57
  58 args.dataset_name = 'a2d2_segmentation_measure' #'tiad_segmentation_infer'   #'cityscapes_segmentation_infer' #'tiad_segmentation'  #'cityscapes_segmentation_measure'
  59
  60 #args.save_path = './data/checkpoints'
  61 args.data_path = '/data/ssd/datasets/a2d2_v2/' #'./data/datasets/cityscapes/data'   #'/data/hdd/datasets/cityscapes_leftImg8bit_sequence_trainvaltest/' #'./data/datasets/cityscapes/data'  #'./data/tiad/data/demoVideo/sequence0021'  #'./data/tiad/data/demoVideo/sequence0025'   #'./data/tiad/data/demoVideo/sequence0001_2017'
  62 #args.pretrained = './data/modelzoo/semantic_segmentation/cityscapes/deeplabv3lite-mobilenetv2/cityscapes_segmentation_deeplabv3lite-mobilenetv2_2019-06-26-08-59-32.pth'
  63 #args.pretrained = './data/checkpoints/tiad_segmentation/2019-10-18_00-50-03_tiad_segmentation_deeplabv3lite_mobilenetv2_ericsun_resize768x384_traincrop768x384_float/checkpoint.pth.tar'
  64
  65 args.pretrained = '/data/files/work/bitbucket_TI/pytorch-jacinto-models/data/checkpoints/a2d2_segmentation/2020-01-25_13-06-18_a2d2_segmentation_deeplabv3lite_mobilenetv2_tv_resize768x384_traincrop768x384/training/model_best_ep172.pth.tar'
  66
  67 args.model_config.input_channels = (3,)
  68 args.model_config.output_type = ['segmentation']
  69 args.model_config.output_channels = None
  70 args.losses = [['segmentation_loss']]
  71 args.metrics = [['segmentation_metrics']]
  72
  73 args.frame_IOU =  False # Print mIOU for each frame
  74 args.shuffle = False
  75
  76 args.num_images = 30000   # Max number of images to run inference on
  77 #args.blend = [False]
  78 #'color', 'blend'
  79 args.viz_op_type = ['color']
  80 args.visualize_gt = False
  81 args.car_mask = False  # False   #True
  82 args.label = [True]    # False   #True
  83 args.label_infer = [True]
  84 args.palette = True
  85 args.start_img_index = 0
  86 args.end_img_index = 0
  87 args.create_video = True  # True   #False
  88 args.depth = [False]
  89
  90 args.epoch_size = 0                     #0 #0.5
  91 args.iter_size = 1                      #2
  92
  93 args.batch_size = 1 #80                  #12 #16 #32 #64
  94 args.img_resize = (384, 768)         #(256,512) #(512,512) # #(1024, 2048) #(512,1024)  #(720, 1280)
  95
  96 args.output_size = (1208, 1920)          #(1024, 2048)
  97 #args.rand_scale = (1.0, 2.0)            #(1.0,2.0) #(1.0,1.5) #(1.0,1.25)
  98
  99 args.quantize = True
 100
 101 args.depth = [False]
 102 args.quantize = False
 103 args.histogram_range = True
 104
 105 #args.image_prenorm=False
 106 #args.image_mean = [0]
 107 #args.image_scale = [1.0]
 108 #args.image_prenorm = False
 109 #args.image_mean = [123.675, 116.28, 103.53]
 110 #args.image_scale = [0.017125, 0.017507, 0.017429]
 111 #args.image_mean = [0]                # image mean for input image normalization
 112 #args.image_scale = [1.0]             # image scaling/mult for input iamge normalization
 113
 114 #save modified files after last commit
 115 #args.save_mod_files = True
 116
 117 args.gpu_mode = True
 118 args.write_layer_ip_op = False
 119 args.generate_onnx = True
 120
 121 ################################
 122 for key in vars(cmds):
 123     if key == 'gpus':
 124         pass # already taken care above, since this has to be done before importing pytorch
 125     elif hasattr(args, key):
 126         value = getattr(cmds, key)
 127         if value != 'None' and value is not None:
 128             setattr(args, key, value)
 129     else:
 130         assert False, f'invalid argument {key}'
 131 #
 132
 133 ################################
 134 #Run the test
 135 infer_pixel2pixel.main(args)
 136