高速公路违停检测
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

53 lines
1.8KB

  1. from models_711.segWaterBuilding import SegModel
  2. from PIL import Image
  3. from torchvision.transforms import transforms
  4. import numpy as np
  5. import cv2
  6. import os
  7. from cv2 import getTickCount, getTickFrequency
  8. import matplotlib.pyplot as plt
  9. def predict_lunkuo(impth=None):
  10. # segmodel = SegModel()
  11. loop_start = getTickCount()
  12. pred = segmodel.eval(image=img)
  13. loop_time = cv2.getTickCount() - loop_start
  14. tool_time = loop_time / (cv2.getTickFrequency())
  15. running_fps = int(1 / tool_time)
  16. print('running_fps:', running_fps)
  17. preds_squeeze = pred.squeeze(0)
  18. preds_squeeze[preds_squeeze != 0] = 255
  19. preds_squeeze = np.array(preds_squeeze.cpu())
  20. preds_squeeze = np.uint8(preds_squeeze)
  21. #print('preds_squeeze:', preds_squeeze.shape)
  22. _, binary = cv2.threshold(preds_squeeze,220,255,cv2.THRESH_BINARY)
  23. contours, hierarchy = cv2.findContours(binary,cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
  24. img_n = cv2.cvtColor(np.asarray(img),cv2.COLOR_RGB2BGR)
  25. img2 = cv2.drawContours(img_n,contours,-1,(0,0,255),8)
  26. # save_path = './' + '00000000000000000000000000001' + '.png'
  27. # cv2.imshow('image',img2)
  28. # cv2.waitKey(0)
  29. plt.figure()
  30. plt.imshow(img2[:,:,[2,1,0]])
  31. # plt.show()
  32. # if __name__ == '__main__':
  33. # impth = "/home/data/lijiwen/wurenjiqifei/images/20211225巡河_10.jpg"
  34. # # to_tensor = transforms.Compose([
  35. # # transforms.ToTensor(),
  36. # # transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
  37. # # ])
  38. # img = Image.open(impth).convert('RGB')
  39. # predict_lunkuo(impth=impth)
  40. if __name__ == '__main__':
  41. impth = '/home/data/lijiwen/wurenjiqifei/bu711/'
  42. segmodel = SegModel()
  43. folders = os.listdir(impth)
  44. for i in range(len(folders)):
  45. imgpath = os.path.join(impth, folders[i])
  46. img = Image.open(imgpath).convert('RGB')
  47. predict_lunkuo(impth=impth)