from PIL import Image
from itertools import groupby
import cv2
import datetime
import os
 
# from core.rabbitmq import MessageQueue
 
THRESHOLD_VALUE = 230  # 二值化時(shí)的閾值
PRETREATMENT_FILE = 'hq'  # 橫切時(shí)臨時(shí)保存的文件夾
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):
  # 獲取區(qū)間
  for k, g in groupby(enumerate(self.continuity_position), lambda x: x[1] - x[0]):
l1 = [j for i, j in g]  # 連續(xù)數(shù)字的列表
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'耗時(shí):{(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))