predict.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. # 环境变量配置,用于控制是否使用GPU
  2. # 说明文档:https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html#gpu
  3. import os
  4. os.environ['CUDA_VISIBLE_DEVICES'] = '0'
  5. import cv2
  6. import numpy as np
  7. import paddlex as pdx
  8. model_dir = 'output/unet_3/best_model'
  9. data_dir = 'google_change_det_dataset'
  10. file_list = 'google_change_det_dataset/val_list.txt'
  11. save_dir = 'output/unet/pred'
  12. if not os.path.exists(save_dir):
  13. os.makedirs(save_dir)
  14. color = [0, 0, 0, 255, 255, 255]
  15. model = pdx.load_model(model_dir)
  16. with open(file_list, 'r') as f:
  17. for line in f:
  18. items = line.strip().split()
  19. img_file_1 = os.path.join(data_dir, items[0])
  20. img_file_2 = os.path.join(data_dir, items[1])
  21. # 原图是tiff格式的图片,PaddleX统一使用gdal库读取
  22. # 因训练数据已经转换成bmp格式,故此处使用opencv读取三通道的tiff图片
  23. #image1 = transforms.Compose.read_img(img_file_1)
  24. #image2 = transforms.Compose.read_img(img_file_2)
  25. image1 = cv2.imread(img_file_1)
  26. image2 = cv2.imread(img_file_2)
  27. image = np.concatenate((image1, image2), axis=-1)
  28. # API说明:https://paddlex.readthedocs.io/zh_CN/develop/apis/models/semantic_segmentation.html#overlap-tile-predict
  29. pred = model.overlap_tile_predict(
  30. img_file=image,
  31. tile_size=(769, 769),
  32. pad_size=[512, 512],
  33. batch_size=4)
  34. pdx.seg.visualize(
  35. img_file_1, pred, weight=0., save_dir=save_dir, color=color)