from PIL import Image
from itertools import groupby
import cv2
import datetime
import os
 
# from core.rabbitmq import MessageQueue
 
THRESHOLD_VALUE = 230  # 二值化时的阈值
PRETREATMENT_FILE = 'hq'  # 横切时临时保存的文件夹
W = 540  # 最小宽度
H = 960  # 最小高度
 
 
class Pretreatment(object):
    __doc__ = "图片横向切割"
 
    def __init__(self, path, save_path, min_size=960):
        self.x = 0
        self.y = 0
        self.img_section = []
        self.continuity_position = []
        self.path = path
        self.save_path = save_path
        self.img_obj = None
        self.min_size = min_size
        self.mkdir(self.save_path)
        self.file_name = self.path.split('/')[-1]
 
    def get_continuity_position_new(self):
        img = cv2.imread(self.path)
        gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, thresh1 = cv2.threshold(gray_image, THRESHOLD_VALUE, 255, cv2.THRESH_BINARY)
 
        width = img.shape[1]
        height = img.shape[0]
        self.x = width
        self.y = height
        for i in range(0, height):
            if thresh1[i].sum() != 255 * width:
                self.continuity_position.append(i)
 
    def filter_rule(self):
        if self.y < self.min_size:
            return True
 
    def mkdir(self, path):
        if not os.path.exists(path):
            os.makedirs(path)
 
    def get_section(self):
        # 获取区间
        for k, g in groupby(enumerate(self.continuity_position), lambda x: x[1] - x[0]):
            l1 = [j for i, j in g]  # 连续数字的列表
            if len(l1) > 1:
                self.img_section.append([min(l1), max(l1)])
 
    def split_img(self):
        print(self.img_section)
        for k, s in enumerate(self.img_section):
            if s:
                if not self.img_obj:
                    self.img_obj = Image.open(self.path)
 
                if self.x < W:
                    return
                if s[1] - s[0] < H:
                    return
                cropped = self.img_obj.crop((0, s[0], self.x, s[1]))  # (left, upper, right, lower)
                self.mkdir(os.path.join(self.save_path, PRETREATMENT_FILE))
                cropped.save(os.path.join(self.save_path, PRETREATMENT_FILE, f"hq_{k}_{self.file_name}"))
 
    def remove_raw_data(self):
        os.remove(self.path)
 
    def main(self):
        # v2
        try:
            self.get_continuity_position_new()
            self.filter_rule()
            self.get_section()
            self.split_img()
        except Exception as e:
            print(self.file_name)
            print(e)
        finally:
            if self.img_obj:
                self.img_obj.close()
 
 
class Longitudinal(Pretreatment):
    def get_continuity_position_new(self):
        print(self.path)
        img = cv2.imread(self.path)
        gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, thresh1 = cv2.threshold(gray_image, THRESHOLD_VALUE, 255, cv2.THRESH_BINARY)
 
        width = img.shape[1]
        height = img.shape[0]
        print(width, height)
        self.x = width
        self.y = height
        for i in range(0, width):
            if thresh1[:, i].sum() != 255 * height:
                self.continuity_position.append(i)
 
    def split_img(self):
        print(self.img_section)
        for k, s in enumerate(self.img_section):
            if s:
                if not self.img_obj:
                    self.img_obj = Image.open(self.path)
                if self.y < H:
                    return
                if s[1] - s[0] < W:
                    return
                cropped = self.img_obj.crop((s[0], 0, s[1], self.y))  # (left, upper, right, lower)
                cropped.save(os.path.join(self.save_path, f"{k}_{self.file_name}"))
 
 
def main(path, save_path):
    starttime = datetime.datetime.now()
    a = Pretreatment(path=path, save_path=save_path)
    a.main()
    for root, dirs, files in os.walk(os.path.join(save_path, PRETREATMENT_FILE)):
        for i in files:
            b = Longitudinal(path=os.path.join(save_path, PRETREATMENT_FILE, i), save_path=save_path)
            b.main()
            os.remove(os.path.join(save_path, PRETREATMENT_FILE, i))
    endtime = datetime.datetime.now()
    print(f'耗时:{(endtime - starttime)}')
 
 
if __name__ == '__main__':
    path = '你图片存放的路径'
    save_path = '要保存的路径'
    for _, _, files in os.walk(path):
        for i in files:
            main(path=os.path.join(path, i), save_path=save_path)
    os.rmdir(os.path.join(save_path, PRETREATMENT_FILE))