#!/usr/bin/python
# -*- coding: utf-8 -*-

import os, sys, time
import numpy, piexif, cv2
from PIL import Image, ImageEnhance
from threading import Thread

#查找灰色区域比例
def find_gray_area(img):
    #灰度化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 灰色颜色范围
    lower_gray = 155
    upper_gray = 180
    # 设置灰色颜色范围
    mask = cv2.inRange(gray, lower_gray, upper_gray)
    # 计算灰色区域比例
    area = cv2.countNonZero(mask)
    return round(area / (img.shape[0] * img.shape[1]), 4)

# 将pillow Image对象转换为opencv的image格式
def convert_pillow_to_opencv(image):
    return cv2.cvtColor(numpy.asarray(image), 4)

# 将opencv的image格式转换为pillow Image对象
def convert_opencv_to_pillow(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return Image.fromarray(image)

#锐化图片
def sharpen_image(image, sigma):
    #高斯模糊
    blurred = cv2.GaussianBlur(image, (0, 0), sigma)
    #计算高斯模糊后的图像与原图像之间的差别
    sharpened = cv2.addWeighted(image, 1.5, blurred, -0.5, 0)
    return sharpened

# 调整图片亮度
def adjust_brightness(image, brightness):
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)
    return image

# gamma correction
def adjust_gamma(img, c=1, g=2.2):
    img = img.astype(float)
    out = img.copy()
    out /= 255.
    out = (1/c * out) ** (1/g)
    out *= 255
    out = out.astype(numpy.uint8)
    return out

# 查找黑色区域融合算法
def fusion_images(image, brightness_value=2.2):
    # 提升亮度将图片充分曝光，使得图片中的黑色部分更加明显，以供后续融合
    boost_brightness_image = cv2.convertScaleAbs(image, alpha=brightness_value)

    # 找出黑色区域图像信息
    lower_black =(0, 0, 0)
    upper_black = (180, 255, 46)

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    mask = cv2.inRange(hsv, lower_black, upper_black)
    inv_mask = cv2.bitwise_not(mask)
    mask = cv2.merge([mask,mask,mask])
    inv_mask = cv2.merge([inv_mask,inv_mask,inv_mask])

    # 融合图像
    boost_brightness_image = cv2.bitwise_and(boost_brightness_image, mask)
    image = cv2.bitwise_and(image,inv_mask)
    image = cv2.add(boost_brightness_image,image)
    return image

def enhanceImages(thread_num):
    while True:
        if len(filelist) == 0:
            break
        else:
            filename = filelist.pop()
        if filename.endswith('_1.jpg'): 
            print('processing ' + filename)
            img = Image.open(os.path.join(path, filename))
            # img = img.rotate(180)
            exif = piexif.load(img.info['exif'])
            exif_bytes = piexif.dump(exif)
            # img = adjust_brightness(img, 1.8)
            img = convert_pillow_to_opencv(img)
            # img = adjust_gamma(img, 1, 2.2)
            img = fusion_images(img, 2.2)
            # img = sharpen_image(img, 5)
            img = convert_opencv_to_pillow(img)
            img.save(os.path.join(path, filename), exif=exif_bytes)
        else:
            gray_area = find_gray_area(cv2.imread(os.path.join(path, filename)))
            print('processing ' + filename + ' gray area: ' + str(gray_area))
            if gray_area > 0.3:
                print('gray area is too large, skip this image')
                os.remove(os.path.join(path, filename))
        

if __name__ == '__main__':
    if len(sys.argv) < 2:
        path = '.'
    else:
        path = sys.argv[1]
    path = os.path.abspath(path)
    filelist = os.listdir(path)

    threads = []
    thread_nums = 10
    for i in range(thread_nums):
        t = Thread(target=enhanceImages, args=(str(i)))
        threads.append(t)
        t.start()


    # for file in os.listdir(path):
    #     if file.endswith('116_1.jpg'):
            
    #         img = Image.open(os.path.join(path, file))
    #         img = img.rotate(180)
    #         exif = piexif.load(img.info['exif'])
    #         exif_bytes = piexif.dump(exif)
    #         # img = adjust_brightness(img, 1.8)
    #         img = convert_pillow_to_opencv(img)
    #         # img = adjust_gamma(img, 1, 2.2)
    #         img = fusion_images(img, 2.2)
    #         # img = sharpen_image(img, 5)
    #         img = convert_opencv_to_pillow(img)
    #         # img.save(os.path.join(path, file), exif=exif_bytes)

    #         img = convert_pillow_to_opencv(img)
    #         win = cv2.namedWindow('image', cv2.WINDOW_NORMAL)
    #         cv2.resizeWindow('image', 1000, 1000)
    #         cv2.imshow('image', img)
            
    #         while True:
    #             if cv2.waitKey(1) & 0xFF == ord('q'):
    #                 break