image_cutout/module/module_rename_pic.py

from UI.rename_pic import Ui_Form as rename_pic_Ui_Form
from PySide6.QtWidgets import *
from PySide6.QtCore import *
# from collections import defaultdict
import os
import log
from PIL import Image, ImageDraw, ImageFont
import time
import settings
# from pyzbar.pyzbar import decode
import shutil
from cv2 import cvtColor, COLOR_RGB2BGR
# cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
import numpy as np
from cv2.wechat_qrcode import WeChatQRCode
from module.module_online_data import GetOnlineData


# pip install opencv-python -i https://pypi.douban.com/simple
# pip install opencv-contrib-python -i https://pypi.douban.com/simple

class Picture:
    def __init__(self, in_path):
        self.im = Image.open(in_path)
        self.x, self.y = self.im.size
        # print(self.x, self.y)

    def save_img(self, outpath, quality=90):
        # self.im = self.im.convert("RGB")
        self.im.save(outpath, quality=quality)

    def resize(self, width):
        re_x = int(width)
        re_y = int(self.y * re_x / self.x)
        self.im = self.im.resize((re_x, re_y), Image.Resampling.LANCZOS)
        self.x, self.y = self.im.size

    def resize_regular(self, width, high):
        self.im = self.im.resize((width, high), Image.Resampling.LANCZOS)

    def corp_square(self):
        if self.y < self.x:
            return
        self.im = self.im.crop((0, int((self.y - self.x) / 2), self.x, self.y - int((self.y - self.x) / 2)))


class RenamePic(QWidget, rename_pic_Ui_Form):
    progress_sign = Signal(dict)
    info_sign = Signal(str)
    text_show = Signal(str)

    def __init__(self, windows=None):
        super().__init__()
        # 加载默认配置
        self.setupUi(self)
        # 0禁用  1进行中  2已结束
        self.state = 2
        self.init()
        self.pic_dir = ""
        self.setFixedSize(self.width(), self.height())
        self.set_state(2)
        self.show()
        QTimer.singleShot(500, self.check_login)

    def check_login(self, *args):
        if not settings.IsLogin:
            a = QMessageBox.question(self, '确认', '请先登录',
                                     QMessageBox.Yes)
            self.close()

    def init(self):
        self.text_show.connect(self.append_text_to_browser)
        self.label_3.mousePressEvent = self.select_pic_dir
        self.label_4.setText("")
        # self.label_6.setText(self.get_len_text(self.to_deal_excel_dir, 50))
        self.pushButton.clicked.connect(self.run)

    def create_folder(self, path):
        if not os.path.exists(path):
            os.makedirs(path)

    def get_len_text(self, text, max_len: int):
        if len(text) > max_len:
            text = text[:int(max_len / 4)] + "..." + text[-1 * int(max_len * 3 / 4):]
        return text

    def select_pic_dir(self, *args):
        pic_dir = QFileDialog.getExistingDirectory(None, "选取文件夹", "")

        if not pic_dir:
            self.label_4.setText("")
            return
        self.pic_dir = pic_dir
        self.label_4.setText(pic_dir)

    def check_path(self, _path):
        if not os.path.exists(_path):
            os.mkdir(_path)
        return True

    def set_state(self, state_value: int):
        # 0禁用  1进行中  2已结束
        if state_value not in [0, 1, 2]:
            return
        self.state = state_value
        if self.state == 0:
            self.progressBar.hide()
            self.pushButton.setEnabled(False)
        if self.state == 1:
            self.progressBar.show()
            self.progressBar.setValue(0)
            self.textBrowser_2.show()
            self.pushButton.setEnabled(False)
            self.textBrowser_2.clear()
        if self.state == 2:
            self.progressBar.hide()
            self.pushButton.setEnabled(True)

    def run(self):
        # 基础检查
        if not settings.IsLogin:
            a = QMessageBox.question(self, '确认', '请先登录',
                                     QMessageBox.Yes)
            return

        self.set_state(1)
        if not self.pic_dir:
            self.text_show.emit("请选择图片文件夹")
            self.set_state(2)
            return
        if self.comboBox.currentText() == "请选择":
            self.text_show.emit("请选择是否将货号添加到图片中")
            self.set_state(2)
            return
        total = 0
        _Type = ['.png', '.PNG', '.jpg', '.JPG', '.gif', '.GIF', '.jpeg', '.JPEG']
        for image_file in os.listdir(self.pic_dir):
            path = "{}\{}".format(self.pic_dir, image_file)
            if os.path.isfile(path):
                if os.path.splitext(image_file)[1] in _Type:
                    total += 1
        if total == 0:
            self.text_show.emit("当前文件夹下没有图片")
            self.set_state(2)
            return
        is_add_text = False
        if self.comboBox.currentText() == "添加":
            is_add_text = True
        self.run_deal(total=total, is_add_text=is_add_text)

        self.set_state(2)

    def send_info(self, data):
        if data["_type"] == "show_p":
            self.progressBar.setValue(data["data"])
        if data["_type"] == "text":
            self.textBrowser_2.append(data["data"])

    def run_deal(self, total, is_add_text=False):
        # 输出目录
        out_put_path = "{}\out_put_{}".format(os.path.split(self.pic_dir)[0], int(time.time()))
        # 无法识别的目录
        error_path = "{}\识别失败".format(out_put_path)
        # 识别成功的目录
        success_path = "{}\识别成功".format(out_put_path)

        self.check_path(out_put_path)
        self.check_path(error_path)
        self.check_path(success_path)

        error_list = []
        _Type = ['.png', '.PNG', '.jpg', '.JPG', '.gif', '.GIF', '.jpeg', '.JPEG']
        last_code = None
        do_n = 0
        for image_file in os.listdir(self.pic_dir):
            if os.path.splitext(image_file)[1] in _Type:
                # 获取当前路径的文件夹名称
                file_path = "{}/{}".format(self.pic_dir, image_file)
                self.send_info({"_type": "text",
                                "data": "{}图片解析".format(image_file)})
                code = self.get_code(file_path)

                if code:
                    if "_" in code:
                        goods_number = code.split("_")[0]
                        numbers_list = [goods_number]
                        r_data = GetOnlineData().get_goods_art_no_info(numbers_list=numbers_list)
                        if goods_number in r_data:
                            code = r_data[goods_number]["商品货号"]
                        else:
                            code = None
                            self.send_info({"_type": "text",
                                            "data": "{}查询不到商品".format(goods_number)})

                if code:
                    # 获取到二维码内容，说明是二维码图片，则不做处理
                    last_code = code
                    continue
                else:
                    # 上个二维码没有解析成功
                    if not last_code:
                        error_list.append(image_file)
                        self.send_info({"_type": "text",
                                        "data": "{}无法识别".format(image_file)})
                        shutil.copyfile(file_path, "{}\{}".format(error_path, image_file))
                    #
                    if last_code:
                        print("{}--------".format(last_code))
                        # 目标文件路径
                        _file_name = "{}{}".format(last_code, os.path.splitext(image_file)[1])
                        dst_file = "{}/{}".format(success_path, _file_name)
                        if os.path.exists(dst_file):
                            self.send_info({"_type": "text",
                                            "data": "{}图片已存在".format(image_file)})
                            last_code = None
                            continue
                        shutil.copy(file_path, dst_file)  # 复制文件
                        pic = Picture(dst_file)
                        pic.resize(width=800)
                        pic.corp_square()  # 居中剪裁
                        if is_add_text:
                            draw = ImageDraw.Draw(pic.im)
                            font_style_1 = ImageFont.truetype(r"ttf\simfang.ttf", 40, encoding="utf-8")
                            draw.text((0, 0), last_code, 0, font=font_style_1)
                        pic.save_img(dst_file)
                        # 处理成功了，清空上个记录
                        last_code = None

            do_n += 1
            self.send_info({"_type": "show_p",
                            "data": int(do_n / total * 100)})
        if error_list:
            self.send_info({"_type": "text",
                            "data": "以下图片无法解析，请自行处理"})
        self.send_info({"_type": "text",
                        "data": "{}".format(error_list)})

        os.startfile(out_put_path)

    def append_text_to_browser(self, text):
        self.textBrowser_2.append(text)

    def create_folder(self, path):
        if not os.path.exists(path):
            os.makedirs(path)
            return False

    def get_code(self, file_path):
        pic = Picture(file_path)
        pic.resize(width=1000)
        im = pic.im
        img = cvtColor(np.asarray(im), COLOR_RGB2BGR)
        # path = r"D:\MyDocuments\PythonCode\MyPython\red_dragonfly\deal_pics\rename_by_qrcode\opencv_3rdparty-wechat_qrcode"
        detector = WeChatQRCode(detector_prototxt_path="qr_mode/detect.prototxt",
                                detector_caffe_model_path="qr_mode/detect.caffemodel",
                                super_resolution_prototxt_path="qr_mode/sr.prototxt",
                                super_resolution_caffe_model_path="qr_mode/sr.caffemodel")

        res, points = detector.detectAndDecode(img)
        if res:
            return res[0]
        return None

    def deal_pic_resize(self, file_path, target_size, times=0, ):
        file_size = int(os.path.getsize(file_path) / 1024)
        if file_size < target_size:
            return
        times += 1
        if times > 2:
            print(file_path, "压缩次数", times)
        k = 0.9
        # if file_size > target_size * 5:
        #     k = 0.7
        pic = Picture(file_path)
        w = int(pic.x * k)
        pic.resize(w)
        pic.save_img(file_path)
        return self.deal_pic_resize(file_path, target_size, times)