import numpy as np
import cv2
# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
"""
Run a Flask REST API exposing one or more YOLOv5s models
"""

import argparse
import io

# import torch
from PIL import Image
from skimage import filters
import os
RESULT_FOLDER = './output_dir/nocutimg' #用于存放检测结果


def draw(img,result,dstDir,filenameList,a):
    """
    为待检测图片 绘制矩形框和结果分数
         Args:
            img (list):利用opesncv按照字节读取的图片
            result (list):检测结果的分数列表
            dstDir (str):检测结果存放位置（./output_dir/nocutimg/exp_1/xxx.jpg）
            filenameList (list):图片名
            a:圆盘内接四边形坐标 （左上角，右上角）

         return:
    """
    try:
         os.makedirs(dstDir)
    except:
         pass
    for i in range(len(filenameList)):
        filename = os.path.join(dstDir,filenameList[i])
        cv2.rectangle(img[i], (a[0], a[1]), (a[2], a[3]), (0, 0, 255), 2)
        cv2.putText(img[i],"score:{}".format(result[i]), (30, 100),cv2.FONT_HERSHEY_PLAIN,7,(0,0,255),2)
        cv2.imwrite(filename,img[i])
    print("检查后图片已存入{}".format(dstDir))

def _imageToMatrix(image):
        """
        根据名称读取图片对象转化矩阵
        :param strName:
        :return: 返回矩阵
        """
        imgMat = np.matrix(image)
        return imgMat

def preImgOps(img):
        """
        图像的预处理操作
        :param img: 图像的而明朝
        :return: 灰度化和resize之后的图片对象
        """
        # 预处理操作
        reImg = cv2.resize(img, (300, 400), interpolation=cv2.INTER_CUBIC)  #
        img2gray = cv2.cvtColor(reImg, cv2.COLOR_BGR2GRAY)  # 将图片压缩为单通道的灰度图
        return reImg,img2gray

def ReadImage(im_file):
    """
        读取图片
        :param im_file: 图像的路径
        :return: 利用opencv读取的完整图像img  托虫板内接四边形图像cutimg   内接四边形的坐标c[xmin,ymin,xmax,ymax] 
    """
    img = cv2.imdecode(np.fromfile(im_file, dtype=np.uint8), cv2.IMREAD_COLOR)
    # imgfile = np.asarray(bytearray(im_file.read()), dtype="uint8")
    # img = cv2.imdecode(imgfile, cv2.IMREAD_ANYCOLOR)
    ImgShape = img.shape
    print("原始图像大小{}".format(ImgShape))
    center = [ImgShape[1]/2,ImgShape[0]/2] #圆盘中心
    print("图像中心位置{}".format(center))
    radius = int(ImgShape[0]/2) #圆盘半径
    print("半径{}".format(radius))
    length = pow(2,0.5)*radius    #圆的内接正方形边长
    # print(length)
    Size = int(length/2)
    xmin = int(center[0] - Size)
    xmax = int(center[0] + Size)
    ymin = int((center[1] - Size)*1.1)
    ymax = int(center[1] + Size)
    c = [xmin,ymin,xmax,ymax]
    # print(c)
    cutimg = img[ymin:ymax,xmin:xmax]
    # img = cv2.imdecode(imgfile, cv2.IMREAD_ANYCOLOR)
    return img,cutimg,c

def getFileName(imageFile):
    url, FileName = os.path.split(imageFile)
    # print(FileName)
    tamp = []
    tamp.append(FileName)
    return tamp

#图像清晰度检测
def predict(imageFile):
    """
    图像清晰度检测
    Args:
        imageFile(dir or file):单张图片或者存放图片的文件夹
    """
    if not os.path.exists(RESULT_FOLDER):
         os.makedirs(RESULT_FOLDER)
    else:
        pass
    dirs = os.listdir(RESULT_FOLDER)
    print(dirs)
    print(imageFile)
    # return str(len(dirs))
    result = [] #存放清晰度评分
    imgList = [] #存放识别后图片
    if os.path.isdir(imageFile):
        imageFileList = [os.path.join(imageFile,x) for x in os.listdir(imageFile)]
        print("imageFileList",imageFileList)
        for image in imageFileList:
            print(image)
            img,cutimg,c = ReadImage(image)
            reImg, img2gray= preImgOps(cutimg)
            f = _imageToMatrix(img2gray)

            tmp = filters.sobel(f)
            source=np.sum(tmp**2)
            source=np.sqrt(source)
            result.append(source)
            imgList.append(img)
        draw(imgList,result,RESULT_FOLDER+'/exp_'+str(len(dirs)+1)+"/", os.listdir(imageFile),c)
    elif os.path.isfile(imageFile):
            print("输入文件")
            img,cutimg,c = ReadImage(imageFile)
            reImg, img2gray= preImgOps(cutimg)
            f = _imageToMatrix(img2gray)

            tmp = filters.sobel(f)
            source=np.sum(tmp**2)
            source=np.sqrt(source)
            result.append(source)
            imgList.append(img)
            draw(imgList,result,RESULT_FOLDER+'/exp_'+str(len(dirs)+1)+"/",getFileName(imageFile),c)
    print(os.path.isdir(imageFile))
    return result

def main():
    """
       imageFile 文件夹或者文件 
    """
    # imageFile = "D:\\新建文件夹 (4)\\WeChat Files\\wxid_uqzo45cgfn1j12\\FileStorage\\File\\2023-06\\图片清晰度评分\\861551058842071-20230612141803.jpg"
    # imageFile = "C:\\Users\\31453\\Desktop\\云飞代码\\imageQuiltyAssign-toyang\\predict\\861551054010806-20230512095242.jpg"
    imageFile = r"E:\emo\output_dir\nocutimg"
    predict(imageFile)

if __name__ == '__main__':
    """
    使用说明：
        修改main()函数中的imageFile：填入待检测图片url or 待检测图片文件夹url
    """
    main()