TuKouKouClient/python/services/SegmentService.py

import os
from middleware import UnicornException


class SegmentService:
    is_fall_file = ["扣图", "已扣图"]
    is_fall_dir = ["微信"]

    def initImages(self, image_list):
        need_cutout_images = []
        for image_path in image_list:
            root_path, file = os.path.split(image_path)
            file_name, file_e = os.path.splitext(file)
            need_cutout_images.append(
                {
                    "file_name": file_name,
                    "file_e": file_e,
                    "file_path": image_path,
                    "file": file,
                    "root_path": root_path,
                    "need_cutout": True,
                }
            )

        for index, image_data in enumerate(need_cutout_images):
            root_path = image_data["root_path"]
            file_name = image_data["file_name"]
            _path_1 = "{}/{}.png".format(root_path, file_name)
            _path_2 = "{}/已扣图/{}.png".format(root_path, file_name)
            if os.path.exists(_path_1):
                need_cutout_images[index]["need_cutout"] = False
                continue
            if os.path.exists(_path_2):
                need_cutout_images[index]["need_cutout"] = False
                continue
        return need_cutout_images
    def check_need_cutout_images(self, root_path):
        _n_c = 0
        need_cutout_images = []
        _Type = [
            ".jpg",
            ".JPG",
            ".jpeg",
            ".JPEG",
            ".png",
            ".PNG",
        ]
        _is_cutout = []
        if not os.path.isdir(root_path):
            raise UnicornException("不是有效路径")
        for file in os.listdir(root_path):
            _n_c += 1
            if _n_c > 500:
                need_cutout_images = []
                raise UnicornException("目录下文件过多，请检查目录是否正确")
            file_path = "{}/{}".format(root_path, file)
            if os.path.isdir(file_path):
                print(file_path)
                if file == "已扣图":
                    # 哪些图片已经有抠图
                    for x_file in os.listdir(file_path):
                        x_file_name, x_file_e = os.path.splitext(x_file)
                        if x_file_e == ".png":
                            continue
                            _is_cutout.append(x_file_name)

        # ===============================================================
        for file in os.listdir(root_path):
            file_path = "{}/{}".format(root_path, file)
            if os.path.isdir(file_path):
                print(file_path)
                # 不是已扣图文件夹，则进行遍历
                f = True
                for i in self.is_fall_dir:
                    if i in file:
                        f = False
                        break
                if f:
                    self.check_need_cutout_images(file_path)

            file_name, file_e = os.path.splitext(file)
            if file_e not in _Type:
                continue

            # 检查文件是否在禁止抠图里
            f = True
            for i in self.is_fall_file:
                if i in file:
                    f = False
                    break
            if not f:
                continue
            need_cutout = False if file_name in _is_cutout else True
            # if os.path.exists("{}/{}.png".format(root_path, file_name)):
            #     need_cutout = False

            # 图片进行处理
            need_cutout_images.append(
                {
                    "file_name": file_name,
                    "file_e": file_e,
                    "file_path": file_path,
                    "file": file,
                    "root_path": root_path,
                    "need_cutout": need_cutout,
                }
            )
        return need_cutout_images