from PIL import Image
import random
import os
import cv2
import shutil


#opencv圖旋轉,並擷取
def rotate_crop(image, angle, center=None, scale=1.0,r = 20): #1
    #旋轉，輸入影象和旋轉角度,r擷取的半徑，center中心，scale縮放
    (h, w) = image.shape[:2] #2
    if center is None: #3
        x = w / 2
        y = h / 2
        center = (x,y) #4 中心
 
    M = cv2.getRotationMatrix2D(center, angle, scale) #5
 
    rotated = cv2.warpAffine(image, M, (w, h)) #6旋轉

    a=int(x-r) # x start
    b=int(x+r) # x end
    c=int(y-r) # y start
    d=int(y+r) # y end
    cropImg = rotated[a:b,c:d]
    return cropImg #7


#旋轉n度，並提取中心小圖
def image_rotate_crop(image_dir,save_dir,n):
    #輸入旋轉角度，影象資料夾地址，儲存的資料夾地址
    for f in os.listdir(image_dir):
        image_path = os.path.join(image_dir, f)
        if os.path.isfile(image_path):
            angle = 0
            im = cv2.imread(image_path)
            while angle < 360:
                re_str = str(angle) + '.jpg'
                name = f.replace('.jpg',re_str)
                save_path = os.path.join(save_dir,name)
                save_image = rotate_crop(im,angle)
                cv2.imwrite(save_path,save_image)
                angle += n

    


 
# 遍歷指定目錄，顯示目錄下的所有檔名
def CropImage4File(filepath,destpath,r):
    pathDir =  os.listdir(filepath) # list all the path or file  in filepath
    for allDir in pathDir:
        child = os.path.join(filepath, allDir)
        dest = os.path.join(destpath,allDir)
        if os.path.isfile(child):
            print(child)
            image = cv2.imread(child)
            sp = image.shape#obtain the image shape
            sz1 = sp[0]#height(rows) of image
            sz2 = sp[1]#width(colums) of image
            #sz3 = sp[2]#the pixels value is made up of three primary colors, here we do not use
            #你想對檔案的操作
            a=int(sz1/2-r) # x start
            b=int(sz1/2+r) # x end
            c=int(sz2/2-r) # y start
            d=int(sz2/2+r) # y end
            cropImg = image[a:b,c:d] #crop the image
            cv2.imwrite(dest,cropImg)# write in destination path             


#讀取檔案內的所有,迭代

def read_file_all(data_dir_path,save_dir):
    for f in os.listdir(data_dir_path):
        data_file_path = os.path.join(data_dir_path, f)
        if os.path.isfile(data_file_path):
            image_rotate(data_file_path,save_dir)
            #print(collected)
        else:
            read_file_all(data_file_path)

def image_rotate(image_path,save_dir):
    #讀取影象,旋轉擴充資料集
    #im = cv2.imread(image_path,0)

    im = Image.open(image_path)
    # 指定逆時針旋轉的角度
    save_path = save_dir + random_name() + '.jpg'
    im.save(save_path)
    #cv2.imwrite(save_path,im)

    #im_rotate1 = im.rotate(90) 
    #save_path = save_dir + random_name() + '.jpg'
    #im_rotate1.save(save_path)
    #cv2.imwrite(save_path,im_rotate1)

    #im_rotate2 = im.rotate(180) 
    #save_path = save_dir + random_name() + '.jpg'
    #im_rotate2.save(save_path)

    #im_rotate3 = im.rotate(270) 
    #save_path = save_dir + random_name() + '.jpg'
    #im_rotate3.save(save_path)

def random_name():
    #隨機數
    a_list = ['0','1','2','3','4','5','6','7','8','9','a','b','c','d','e']
    name = random.sample(a_list,8)
    file_name = "".join(name)
    return file_name

 


def image_class(image_dir,save_dir):
    #根據檔名對影象進行分類
    for f in os.listdir(image_dir):
        image_path = os.path.join(image_dir, f)
        if os.path.isfile(image_path):
            class_name = f.split(')')[0]
            save_class_dir = save_dir + '/' + str(class_name)
            if os.path.exists(save_class_dir):
                save_class_path = os.path.join(save_class_dir, f)
                try:
                    shutil.copyfile(image_path,save_class_path)
                except:
                    pass
            else:
                os.mkdir(save_class_dir)
                save_class_path = os.path.join(save_class_dir, f)
                try:
                    shutil.copyfile(image_path,save_class_path)
                except:
                    pass



if __name__ == '__main__':
    # data_dir_path = 'E:/pcb_image_data/data_2500/test_2500/True' #讀取的檔案路徑
    # save_dir = 'E:/pcb_image_data/data_2500_more/test_2500/True/'#儲存的檔案路徑
    # read_file_all(data_dir_path,save_dir)

    # 擷取小圖
    filepath = 'E:/pcb_defect_image/class_image/Type36/True'  # source images
    destpath = 'E:/pcb_defect_image/class_image/small_image/40_40_36/True'  # resized images saved here
    CropImage4File(filepath, destpath, 20) #半徑

    # 旋轉
    # scr_dir = 'E:/pcb_image_data/2018_11_7_h/train_20/20'
    # dst_dir = 'E:/pcb_image_data/2018_11_7_h/train_20/True/'
    # read_file_all(scr_dir,dst_dir)

    # 旋轉並擷取小圖
    #image_dir = 'E:/pcb_defect_image/2018_11_26/train/True'
    #save_dir = 'E:/pcb_defect_image/2018_11_26/image_13/train/True'

    #image_rotate_crop(image_dir, save_dir, 30)

    #根據檔名對圖片分類
    #scr_dir = 'E:/pcb_image_data/2018_11_8_h_hu/ALL/true'
    #dst_dir = 'E:/pcb_image_data/2018_11_8_h_hu/ALL/true_class'
    #image_class(scr_dir,dst_dir)