utils.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. import numpy as np
  2. import cv2
  3. # YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
  4. """
  5. Run a Flask REST API exposing one or more YOLOv5s models
  6. """
  7. import argparse
  8. import io
  9. # import torch
  10. from PIL import Image
  11. from skimage import filters
  12. import os
  13. RESULT_FOLDER = './output_dir/nocutimg' #用于存放检测结果
  14. def draw(img,result,dstDir,filenameList,a):
  15. """
  16. 为待检测图片 绘制矩形框和结果分数
  17. Args:
  18. img (list):利用opesncv按照字节读取的图片
  19. result (list):检测结果的分数列表
  20. dstDir (str):检测结果存放位置(./output_dir/nocutimg/exp_1/xxx.jpg)
  21. filenameList (list):图片名
  22. a:圆盘内接四边形坐标 (左上角,右上角)
  23. return:
  24. """
  25. try:
  26. os.makedirs(dstDir)
  27. except:
  28. pass
  29. for i in range(len(filenameList)):
  30. filename = os.path.join(dstDir,filenameList[i])
  31. cv2.rectangle(img[i], (a[0], a[1]), (a[2], a[3]), (0, 0, 255), 2)
  32. cv2.putText(img[i],"score:{}".format(result[i]), (30, 100),cv2.FONT_HERSHEY_PLAIN,7,(0,0,255),2)
  33. cv2.imwrite(filename,img[i])
  34. print("检查后图片已存入{}".format(dstDir))
  35. def _imageToMatrix(image):
  36. """
  37. 根据名称读取图片对象转化矩阵
  38. :param strName:
  39. :return: 返回矩阵
  40. """
  41. imgMat = np.matrix(image)
  42. return imgMat
  43. def preImgOps(img):
  44. """
  45. 图像的预处理操作
  46. :param img: 图像的而明朝
  47. :return: 灰度化和resize之后的图片对象
  48. """
  49. # 预处理操作
  50. reImg = cv2.resize(img, (300, 400), interpolation=cv2.INTER_CUBIC) #
  51. img2gray = cv2.cvtColor(reImg, cv2.COLOR_BGR2GRAY) # 将图片压缩为单通道的灰度图
  52. return reImg,img2gray
  53. def ReadImage(im_file):
  54. """
  55. 读取图片
  56. :param im_file: 图像的路径
  57. :return: 利用opencv读取的完整图像img 托虫板内接四边形图像cutimg 内接四边形的坐标c[xmin,ymin,xmax,ymax]
  58. """
  59. img = cv2.imdecode(np.fromfile(im_file, dtype=np.uint8), cv2.IMREAD_COLOR)
  60. # imgfile = np.asarray(bytearray(im_file.read()), dtype="uint8")
  61. # img = cv2.imdecode(imgfile, cv2.IMREAD_ANYCOLOR)
  62. ImgShape = img.shape
  63. print("原始图像大小{}".format(ImgShape))
  64. center = [ImgShape[1]/2,ImgShape[0]/2] #圆盘中心
  65. print("图像中心位置{}".format(center))
  66. radius = int(ImgShape[0]/2) #圆盘半径
  67. print("半径{}".format(radius))
  68. length = pow(2,0.5)*radius #圆的内接正方形边长
  69. # print(length)
  70. Size = int(length/2)
  71. xmin = int(center[0] - Size)
  72. xmax = int(center[0] + Size)
  73. ymin = int((center[1] - Size)*1.1)
  74. ymax = int(center[1] + Size)
  75. c = [xmin,ymin,xmax,ymax]
  76. # print(c)
  77. cutimg = img[ymin:ymax,xmin:xmax]
  78. # img = cv2.imdecode(imgfile, cv2.IMREAD_ANYCOLOR)
  79. return img,cutimg,c
  80. def getFileName(imageFile):
  81. url, FileName = os.path.split(imageFile)
  82. # print(FileName)
  83. tamp = []
  84. tamp.append(FileName)
  85. return tamp
  86. #图像清晰度检测
  87. def predict(imageFile):
  88. """
  89. 图像清晰度检测
  90. Args:
  91. imageFile(dir or file):单张图片或者存放图片的文件夹
  92. """
  93. if not os.path.exists(RESULT_FOLDER):
  94. os.makedirs(RESULT_FOLDER)
  95. else:
  96. pass
  97. dirs = os.listdir(RESULT_FOLDER)
  98. print(dirs)
  99. print(imageFile)
  100. # return str(len(dirs))
  101. result = [] #存放清晰度评分
  102. imgList = [] #存放识别后图片
  103. if os.path.isdir(imageFile):
  104. imageFileList = [os.path.join(imageFile,x) for x in os.listdir(imageFile)]
  105. print("imageFileList",imageFileList)
  106. for image in imageFileList:
  107. print(image)
  108. img,cutimg,c = ReadImage(image)
  109. reImg, img2gray= preImgOps(cutimg)
  110. f = _imageToMatrix(img2gray)
  111. tmp = filters.sobel(f)
  112. source=np.sum(tmp**2)
  113. source=np.sqrt(source)
  114. result.append(source)
  115. imgList.append(img)
  116. draw(imgList,result,RESULT_FOLDER+'/exp_'+str(len(dirs)+1)+"/", os.listdir(imageFile),c)
  117. elif os.path.isfile(imageFile):
  118. print("输入文件")
  119. img,cutimg,c = ReadImage(imageFile)
  120. reImg, img2gray= preImgOps(cutimg)
  121. f = _imageToMatrix(img2gray)
  122. tmp = filters.sobel(f)
  123. source=np.sum(tmp**2)
  124. source=np.sqrt(source)
  125. result.append(source)
  126. imgList.append(img)
  127. draw(imgList,result,RESULT_FOLDER+'/exp_'+str(len(dirs)+1)+"/",getFileName(imageFile),c)
  128. print(os.path.isdir(imageFile))
  129. return result
  130. def main():
  131. """
  132. imageFile 文件夹或者文件
  133. """
  134. # imageFile = "D:\\新建文件夹 (4)\\WeChat Files\\wxid_uqzo45cgfn1j12\\FileStorage\\File\\2023-06\\图片清晰度评分\\861551058842071-20230612141803.jpg"
  135. # imageFile = "C:\\Users\\31453\\Desktop\\云飞代码\\imageQuiltyAssign-toyang\\predict\\861551054010806-20230512095242.jpg"
  136. imageFile = r"E:\emo\output_dir\nocutimg"
  137. predict(imageFile)
  138. if __name__ == '__main__':
  139. """
  140. 使用说明:
  141. 修改main()函数中的imageFile:填入待检测图片url or 待检测图片文件夹url
  142. """
  143. main()