def Generate(args): import PIL.Image as Image import matplotlib.pyplot as plt #model = Unet(3, 3) #vgg_model = VGGNet(requires_grad=True, remove_fc=True) #model = FCN8s(pretrained_net=vgg_model, n_class=3) model = ESPNet_Encoder(3, p=2, q=3) model.load_state_dict(torch.load(args.ckpt,map_location='cpu')) model.eval() count=0 root = './data_road/training/' print("Generate...") for filename in os.listdir(root): if filename == 'image': for filename2 in os.listdir(root+filename): if filename2 == '.ipynb_checkpoints': break imgroot=os.path.join(root+'image/'+filename2) name = filename2[:-4] print("Image Processing: ",name) img = Image.open(imgroot) img = x_transforms(img) img = img.view(1,3,352,1216) #for Fcn with torch.no_grad(): output= model(img) output = torch.softmax(output,dim=1) N, _, h, w = output.shape pred = output.transpose(0, 2).transpose(3, 1).reshape(-1, 3).argmax(axis=1).reshape(N, h, w) #class 3 pred = pred.squeeze(0) Decode_image(pred,name)
def main(): #args = parse_args() '''model = init_detector( args.config, args.checkpoint, device=torch.device('cuda', args.device)) ''' model = init_detector( '../configs/cityscapes/mask_rcnn_r50_fpn_1x_cityscapes.py', '../checkpoints/mask_rcnn_r50_fpn_1x_city_20190727-9b3c56a5.pth', device=torch.device('cuda', 0)) modelseg = ESPNet_Encoder(3, p=2, q=3).to(device) modelseg.load_state_dict( torch.load('bdd_weights_20_ESPNET_road.pth', map_location='cpu')) modelseg.eval() #camera = cv2.VideoCapture(args.camera_id) camera = cv2.VideoCapture('umgt.avi') print('Press "Esc", "q" or "Q" to exit.') if camera.isOpened(): while True: ret_val, img = camera.read() imgroad = RoadSeg(img, modelseg) #imgroad = evaluateModel(img,modelseg) result = inference_detector(model, img) ch = cv2.waitKey(1) if ch == 27 or ch == ord('q') or ch == ord('Q'): break #show_result(img, result, model.CLASSES, score_thr=args.score_thr, wait_time=1) show_result(imgroad, result, model.CLASSES, score_thr=0.5, wait_time=1) cv2.destroyAllWindows()
def oonx(model): dummy_input = torch.randn(1, 3, 720, 1280).to(device) model = ESPNet_Encoder(3, p=2, q=3).to(device) model.load_state_dict(torch.load(args.ckpt, map_location='cpu')) model.eval() torch.onnx.export(model, dummy_input, "bbd.onnx", verbose=True)