CameraMachine/python/service/generate_main_image/module_generate_main_image.py

"""
步骤：
1、抠图
2、色阶处理，保留阴影
3、定位商品中心，并进行抠图
4、缩放定位，并粘贴到中心位置
5、粘贴水印
"""
import copy
import time
import cv2
import numpy as np
from PIL import Image, ImageEnhance
import os
from module.log.log import MyLogger

if __name__ != '__main__':
    import settings


class GeneratePic(object):
    def __init__(self):
        self.logger = MyLogger()
        pass

    def to_resize(self, _im, width=None, high=None):
        _im_x, _im_y = _im.size
        if width and high:
            if _im_x >= _im_y:
                high = None
            else:
                width = None
        if width:
            re_x = int(width)
            re_y = int(_im_y * re_x / _im_x)
        else:
            re_y = int(high)
            re_x = int(_im_x * re_y / _im_y)
        _im = _im.resize((re_x, re_y))
        return _im

    def clean_colors(self, img):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return img

    def levels_adjust(self, img, Shadow, Midtones, Highlight, OutShadow, OutHighlight, Dim):
        # print(img)
        # dim = 3的时候调节RGB三个分量， 0调节B，1调节G，2调节R
        if Dim == 3:
            mask_shadow = img < Shadow
            img[mask_shadow] = Shadow
            mask_Highlight = img > Highlight
            img[mask_Highlight] = Highlight
        else:
            mask_shadow = img[..., Dim] < Shadow
            img[mask_shadow] = Shadow
            mask_Highlight = img[..., Dim] > Highlight
            img[mask_Highlight] = Highlight

        if Dim == 3:
            Diff = Highlight - Shadow
            rgbDiff = img - Shadow
            clRgb = np.power(rgbDiff / Diff, 1 / Midtones)
            outClRgb = clRgb * (OutHighlight - OutShadow) / 255 + OutShadow
            data = np.array(outClRgb * 255, dtype='uint8')
            img = data
        else:
            Diff = Highlight - Shadow
            rgbDiff = img[..., Dim] - Shadow
            clRgb = np.power(rgbDiff / Diff, 1 / Midtones)
            outClRgb = clRgb * (OutHighlight - OutShadow) / 255 + OutShadow
            data = np.array(outClRgb * 255, dtype='uint8')
            img[..., Dim] = data
        return img

    def calculate_average_brightness_opencv(self, img_gray, rows_to_check):
        # 二值化的图片 CV对象
        height, width = img_gray.shape

        brightness_list = []
        for row in rows_to_check:
            if 0 <= row < height:
                # 直接计算该行的平均亮度
                row_data = img_gray[row, :]
                average_brightness = np.mean(row_data)
                brightness_list.append(average_brightness)
            else:
                print(f"警告：行号{row}超出图片范围，已跳过。")

        return brightness_list

    def run(self, image_path, cut_image_path, out_path, image_deal_mode=0, image_index=99,
            out_pic_size=1024, is_logo=True, out_process_path_1=None, out_process_path_2=None,
            resize_mode=None, **kwargs):  # im 为cv对象
        """
        image_path:原始图
        cut_image_path:抠图结果
        out_path:输出文件夹
        image_deal_mode:图片处理模式，1表示需要镜像处理
        """

        # ==========先进行剪切原图
        im = Image.open(image_path)
        old_x, old_y = im.size

        cut_image = Image.open(cut_image_path)
        try:
            x1, y1, x2, y2 = cut_image.getbbox()
            # im_shadow.crop((x1, y1, x2, y2)).show()# 切图
        except BaseException as e:
            print("抠图失败")
            return False

        goods_w, goods_h = x2 - x1, y2 - y1
        _w, _h = int(goods_w / 10), int(goods_h / 10)  # 上下左右扩展位置
        new_x1, new_y1, new_x2, new_y2 = x1 - _w, y1 - _h, x2 + _w, y2 + _h  # 防止超限
        new_x1 = 0 if new_x1 < 0 else new_x1
        new_y1 = 0 if new_y1 < 0 else new_y1
        new_x2 = old_x if new_x2 > old_x else new_x2
        new_y2 = old_y if new_y2 > old_y else new_y2

        im = im.crop((new_x1, new_y1, new_x2, new_y2))  # 切图
        cut_image = cut_image.crop((new_x1, new_y1, new_x2, new_y2))  # 切图

        new_x, new_y = im.size
        # ================自动色阶处理

        _im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
        # 背景阴影
        im_shadow = cv2.cvtColor(_im, cv2.COLOR_BGR2GRAY)

        s = time.time()
        rows = [int(new_y * 19.8 / 20)]

        # draw = ImageDraw.Draw(im)
        # for row in rows:
        #     draw.line((0, row, new_x, row), fill=128, width=3)
        # im.show()

        _im_shadow = copy.copy(im_shadow)
        Midtones = 0.52
        Highlight = 235
        k = 20
        while k:
            k -= 1
            Midtones += 0.1
            Highlight -= 5
            _im_shadow = self.levels_adjust(img=im_shadow, Shadow=0, Midtones=Midtones, Highlight=Highlight,
                                            OutShadow=0,
                                            OutHighlight=255, Dim=3)
            brightness_list = self.calculate_average_brightness_opencv(img_gray=_im_shadow, rows_to_check=rows)
            print(brightness_list)
            if brightness_list[0] >= 254:
                break

        print(time.time() - s)

        im_shadow = Image.fromarray(cv2.cvtColor(_im_shadow, cv2.COLOR_BGR2RGB))
        im_shadow.paste(cut_image, (0, 0), cut_image)  # 把原图粘贴回去，避免色差

        # 保存带有阴影的底图，没有logo
        if out_process_path_1:
            out_image_1 = im_shadow.copy()
            if image_deal_mode == 1:
                out_image_1 = out_image_1.transpose(Image.FLIP_LEFT_RIGHT)
            out_image_1.save(out_process_path_1)

        # 保存抠图结果，没有底图，没有logo
        if out_process_path_2:
            out_image_2 = cut_image.copy()
            if image_deal_mode == 1:
                out_image_2 = out_image_2.transpose(Image.FLIP_LEFT_RIGHT)
            out_image_2.save(out_process_path_2)

        # im_shadow.show()
        # 主图物体的缩放依据大小
        if resize_mode is None:
            im_shadow = self.to_resize(_im=im_shadow, width=1400, high=1400)
            cut_image = self.to_resize(_im=cut_image, width=1400, high=1400)

        elif resize_mode == 1:
            im_shadow = self.to_resize(_im=im_shadow, width=1400, high=1400)
            cut_image = self.to_resize(_im=cut_image, width=1400, high=1400)

        elif resize_mode == 2:
            if im_shadow.height <= im_shadow.width * 1.2:
                im_shadow = self.to_resize(_im=im_shadow, width=650)
                cut_image = self.to_resize(_im=cut_image, width=650)
            else:
                im_shadow = self.to_resize(_im=im_shadow, high=1400)
                cut_image = self.to_resize(_im=cut_image, high=1400)

        if image_deal_mode == 1:
            # 翻转
            im_shadow = im_shadow.transpose(Image.FLIP_LEFT_RIGHT)
            cut_image = cut_image.transpose(Image.FLIP_LEFT_RIGHT)

        # 创建底层背景
        image_bg = Image.new("RGB", (1600, 1600), (255, 255, 255))

        image_bg_x, image_bg_y = image_bg.size
        image_x, image_y = im_shadow.size

        _x = int((image_bg_x - image_x) / 2)
        _y = int((image_bg_y - image_y) / 2)

        image_bg.paste(im_shadow, (_x, _y))
        image_bg.paste(cut_image, (_x, _y), cut_image)  # 再叠加原图避免色差

        if is_logo:
            logo_path = ""
            if settings.PROJECT == "红蜻蜓":
                logo_path = r"resources\LOGO\HQT\logo.png"
            elif settings.PROJECT == "惠利玛":
                if "小苏" in settings.Company:
                    logo_path = r"resources\LOGO\xiaosushuoxie\logo.png"
                elif "惠利玛" in settings.Company:
                    logo_path = r"resources\LOGO\HLM\logo.png"
                else:
                    pass

            if not logo_path:
                logo_im = Image.new("RGBA", (1600, 1600), (0, 0, 0, 0))
            else:
                if os.path.exists(logo_path):
                    logo_im = Image.open(logo_path)
                else:
                    logo_im = Image.new("RGBA", (1600, 1600), (0, 0, 0, 0))

            image_bg.paste(logo_im, (0, 0), logo_im)


        image_bg = image_bg.resize((out_pic_size, out_pic_size), Image.BICUBIC)

        # image_bg.show()
        image_bg.save(out_path, quality=100, dpi=(300, 300), format="JPEG")
        return True


class GeneratePicPiJu(object):
    def __init__(self):
        pass

    def calculate_average_brightness_opencv(self, img_gray, rows_to_check):
        # 二值化的图片 CV对象
        height, width = img_gray.shape

        brightness_list = []
        for row in rows_to_check:
            if 0 <= row < height:
                # 直接计算该行的平均亮度
                row_data = img_gray[row, :]
                average_brightness = np.mean(row_data)
                brightness_list.append(average_brightness)
            else:
                print(f"警告：行号{row}超出图片范围，已跳过。")

        return brightness_list

    def to_resize(self, _im, width=None, high=None):
        _im_x, _im_y = _im.size
        if width and high:
            if _im_x >= _im_y:
                high = None
            else:
                width = None
        if width:
            re_x = int(width)
            re_y = int(_im_y * re_x / _im_x)
        else:
            re_y = int(high)
            re_x = int(_im_x * re_y / _im_y)
        _im = _im.resize((re_x, re_y))
        return _im

    def run(self, image_path, cut_image_path, out_path, image_deal_mode=0, image_index=99,
            out_pic_size=800, is_logo=True, out_process_path_1=None, out_process_path_2=None, max_box=None):  # im 为cv对象
        """
        image_path:原始图
        cut_image_path:抠图结果
        out_path:输出文件夹
        image_deal_mode:图片处理模式，1表示需要镜像处理
        """

        # ==========先进行剪切原图
        im = Image.open(image_path)
        old_x, old_y = im.size

        cut_image = Image.open(cut_image_path)
        try:
            x1, y1, x2, y2 = cut_image.getbbox()
            # im_shadow.crop((x1, y1, x2, y2)).show()# 切图
        except BaseException as e:
            print("抠图失败")
            return False

        goods_w, goods_h = x2 - x1, y2 - y1
        _w, _h = int(goods_w / 10), int(goods_h / 10)  # 上下左右扩展位置
        new_x1, new_y1, new_x2, new_y2 = x1 - _w, y1 - _h, x2 + _w, y2 + _h  # 防止超限
        new_x1 = 0 if new_x1 < 0 else new_x1
        new_y1 = 0 if new_y1 < 0 else new_y1
        new_x2 = old_x if new_x2 > old_x else new_x2
        new_y2 = old_y if new_y2 > old_y else new_y2

        im = im.crop((new_x1, new_y1, new_x2, new_y2))  # 切图
        cut_image = cut_image.crop((new_x1, new_y1, new_x2, new_y2))  # 切图

        new_x, new_y = im.size
        # ================自动色阶处理

        _im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
        # 背景阴影
        im_shadow = cv2.cvtColor(_im, cv2.COLOR_BGR2GRAY)

        s = time.time()
        rows = [int(new_y * 19.8 / 20)]

        # draw = ImageDraw.Draw(im)
        # for row in rows:
        #     draw.line((0, row, new_x, row), fill=128, width=3)
        # im.show()

        _im_shadow = copy.copy(im_shadow)
        Midtones = 0.52
        Highlight = 235
        k = 20
        while k:
            k -= 1
            Midtones += 0.1
            Highlight -= 5
            _im_shadow = self.levels_adjust(img=im_shadow, Shadow=0, Midtones=Midtones, Highlight=Highlight,
                                            OutShadow=0,
                                            OutHighlight=255, Dim=3)
            brightness_list = self.calculate_average_brightness_opencv(img_gray=_im_shadow, rows_to_check=rows)
            print(brightness_list)
            if brightness_list[0] >= 254:
                break

        print(time.time() - s)

        im_shadow = Image.fromarray(cv2.cvtColor(_im_shadow, cv2.COLOR_BGR2RGB))
        im_shadow.paste(cut_image, (0, 0), cut_image)  # 把原图粘贴回去，避免色差

        # 保存带有阴影的底图，没有logo
        if out_process_path_1:
            out_image_1 = im_shadow.copy()
            if image_deal_mode == 1:
                out_image_1 = out_image_1.transpose(Image.FLIP_LEFT_RIGHT)
            out_image_1.save(out_process_path_1)

        # 保存抠图结果，没有底图，没有logo
        if out_process_path_2:
            out_image_2 = cut_image.copy()
            if image_deal_mode == 1:
                out_image_2 = out_image_2.transpose(Image.FLIP_LEFT_RIGHT)
            out_image_2.save(out_process_path_2)

        # im_shadow.show()
        if max_box:
            im_shadow = self.to_resize(_im=im_shadow, width=max_box[0], high=max_box[1])
            cut_image = self.to_resize(_im=cut_image, width=max_box[0], high=max_box[1])
        else:
            if image_index != 3:
                im_shadow = self.to_resize(_im=im_shadow, width=1000, high=1000)
                cut_image = self.to_resize(_im=cut_image, width=1000, high=1000)
            else:
                im_shadow = self.to_resize(_im=im_shadow, width=650)
                cut_image = self.to_resize(_im=cut_image, width=650)

        if image_deal_mode == 1:
            # 翻转
            im_shadow = im_shadow.transpose(Image.FLIP_LEFT_RIGHT)
            cut_image = cut_image.transpose(Image.FLIP_LEFT_RIGHT)

        # 创建底层背景
        image_bg = Image.new("RGB", (1600, 1600), (255, 255, 255))

        image_bg_x, image_bg_y = image_bg.size
        image_x, image_y = im_shadow.size

        _x = int((image_bg_x - image_x) / 2)
        _y = int((image_bg_y - image_y) / 2)

        image_bg.paste(im_shadow, (_x, _y))
        image_bg.paste(cut_image, (_x, _y), cut_image)

        if is_logo:
            logo_path = ""

            if settings.PROJECT == "红蜻蜓":
                logo_path = r"resources\LOGO\HQT\logo.png"
            elif settings.PROJECT == "惠利玛":
                if "小苏" in settings.Company:
                    logo_path = r"resources\LOGO\xiaosushuoxie\logo.png"
                elif "惠利玛" in settings.Company:
                    logo_path = r"resources\LOGO\HLM\logo.png"
                else:
                    pass

            if not logo_path:
                logo_im = Image.new("RGBA", (1600, 1600), (0, 0, 0, 0))
            else:
                if os.path.exists(logo_path):
                    logo_im = Image.open(logo_path)
                else:
                    logo_im = Image.new("RGBA", (1600, 1600), (0, 0, 0, 0))

            image_bg.paste(logo_im, (0, 0), logo_im)

        image_bg = image_bg.resize((out_pic_size, out_pic_size), Image.BICUBIC)

        # image_bg.show()
        image_bg.save(out_path, quality=100, dpi=(300, 300), format="JPEG")
        return True

    def increase_brightness(self, im, brightness_factor):
        # 打开图片
        # 将图片转换为RGB模式
        # image = image.convert("RGBA")
        # 获取图片的亮度值
        enhancer = ImageEnhance.Brightness(im)
        brightness = enhancer.enhance(brightness_factor)
        return brightness

    def clean_colors(self, img):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return img

    def levels_adjust(self, img, Shadow, Midtones, Highlight, OutShadow, OutHighlight, Dim):
        # print(img)
        # dim = 3的时候调节RGB三个分量， 0调节B，1调节G，2调节R
        if Dim == 3:
            mask_shadow = img < Shadow
            img[mask_shadow] = Shadow
            mask_Highlight = img > Highlight
            img[mask_Highlight] = Highlight
        else:
            mask_shadow = img[..., Dim] < Shadow
            img[mask_shadow] = Shadow
            mask_Highlight = img[..., Dim] > Highlight
            img[mask_Highlight] = Highlight

        if Dim == 3:
            Diff = Highlight - Shadow
            rgbDiff = img - Shadow
            clRgb = np.power(rgbDiff / Diff, 1 / Midtones)
            outClRgb = clRgb * (OutHighlight - OutShadow) / 255 + OutShadow
            data = np.array(outClRgb * 255, dtype='uint8')
            img = data
        else:
            Diff = Highlight - Shadow
            rgbDiff = img[..., Dim] - Shadow
            clRgb = np.power(rgbDiff / Diff, 1 / Midtones)
            outClRgb = clRgb * (OutHighlight - OutShadow) / 255 + OutShadow
            data = np.array(outClRgb * 255, dtype='uint8')
            img[..., Dim] = data
        return img


if __name__ == '__main__':
    # 测试皮具主图生成
    for jpg_name in os.listdir("image"):
        if "png" in jpg_name:
            continue
        file_name = os.path.splitext(jpg_name)[0]
        original_image_path = r"image\{}.jpg".format(file_name)
        original_move_bg_image_path = r"image\{}.png".format(file_name)
        out_path = r"out\{}.jpg".format(file_name)
        is_image_deal_mode = 1
        resize_mode = 1

        print(original_image_path)
        GeneratePicPiJu().run(image_path=original_image_path,
                              cut_image_path=original_move_bg_image_path,
                              out_path=out_path,
                              image_deal_mode=is_image_deal_mode,
                              # resize_mode=resize_mode,
                              out_pic_size=1600,
                              is_logo=False,
                              max_box=None
                              )