image_cutout/module/remove_bg_pixian.py

import copy
import os
from PIL import Image
from module.other.remove_bg_ali import RemoveBgALi
import requests
from io import BytesIO


class Segment(object):
    def __init__(self):
        self.k = "pxnib99dbchtmdm"
        self.s = "ub9uj5678gs4m2bnrass1t3tn6ughlk065ianosk06akagolcr2u"

    def get_no_bg_goods2(self, file_path=None, _im=None):
        im = _im
        img = BytesIO()
        try:
            im.save(img, format='JPEG')  # format: PNG or JPEG
        except:
            im.save(img, format='PNG')  # format: PNG or JPEG
        img.seek(0)  # rewind to the start

        # img = "https://ossimg.valimart.net/uploads/vali_ai/20241011/172864207036098.png"

        response = requests.post(
            'http://47.76.110.118:27777/api/v2/remove-background',
            files={'image.url': img},
            data={
                # TODO: Add more upload options here
            },
            auth=(self.k, self.s)
        )
        # print(response.content)
        if response.status_code == requests.codes.ok:
            return response.content, ""
        else:
            return None, response.content

    def get_no_bg_goods_by_url(self, url):
        response = requests.post(
            'https://api.pixian.ai/api/v2/remove-background',
            data={
                'image.url': url
            },
            auth=(self.k, self.s)
        )

        if response.status_code == requests.codes.ok:
            return response.content, ""
        else:
            print("response.status_code:", response.status_code)
            return None, response.content

    def get_no_bg_goods(self, file_path=None, _im=None, key=None):
        im = _im
        img = BytesIO()
        try:
            im.save(img, format='JPEG')  # format: PNG or JPEG
        except:
            im.save(img, format='PNG')  # format: PNG or JPEG
        img.seek(0)  # rewind to the start

        if key:
            # todo 切换key
            auth = key
            # auth = (self.k, self.s)
        else:
            auth = (self.k, self.s)


        try:
            response = requests.post(
                'https://api.pixian.ai/api/v2/remove-background',
                files={'image': img},
                data={
                    # TODO: Add more upload options here
                },
                auth=auth,
                timeout=40
            )
        except BaseException as e:
            data = {"im": None,
                    "status_code": "time_out",
                    "message":"{}".format(e)
                    }
            return data

        # print(response.content)

        data = {"im": None,
                "status_code": response.status_code, }

        if response.status_code == requests.codes.ok:
            data["im"] = Image.open(BytesIO(response.content))

        return data


class Picture:
    def __init__(self, in_path, im=None):
        if im:
            self.im = im
        else:
            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.BICUBIC)
        self.x, self.y = self.im.size

    def resize_by_heigh(self, heigh):
        re_y = int(heigh)
        re_x = int(self.x * re_y / self.y)
        self.im = self.im.resize((re_x, re_y), Image.BICUBIC)
        self.x, self.y = self.im.size


class RemoveBgPiXian(object):
    def __init__(self):
        self.segment = Segment()
        self.r = RemoveBgALi()

    def direct_matting_image(self, image):
        # todo 不能超过32,000,000尺寸的数据
        x, y = image.size
        f = False
        if x * y > 32000000:
            r = 32000000 / x * y
            image = image.resize(size=(int(x * r), int(y * r)))
            f = True
        pic, _ = self.segment.get_no_bg_goods(file_path=None, _im=image)
        if not pic:
            return None, _
        _img_im = Image.open(BytesIO(pic))  # 阿里返回的抠图结果 已转PIL对象
        if f:
            _img_im = _img_im.resize(size=(x, y))
        return _img_im, ""

    def run_by_image_url(self, url):
        pic, _ = self.segment.get_no_bg_goods_by_url(url)
        if pic is not None:
            _img_im = Image.open(BytesIO(pic))
            return _img_im, None
        else:
            return None, _

    def run_by_image_im(self, im, key):
        return self.segment.get_no_bg_goods(_im=im, key=key)

    def get_image_cut(self, file_path, out_file_path=None, original_im=None, image_preprocessing=False, is_test=False):
        if original_im:
            original_pic = Picture(in_path=None, im=original_im)
        else:
            original_pic = Picture(file_path)

        if original_pic.im.mode != "RGB":
            print("抠图图片不能是PNG")
            return False, {"data": "抠图图片不能是PNG"}

        if is_test:
            cut_image = self.r.get_image_cut(file_path=None, out_file_path=None, original_im=original_pic.im)
            if out_file_path:
                cut_image.save(out_file_path)
            return True, {}

        if image_preprocessing:
            cut_image = self.r.get_image_cut(file_path=None, out_file_path=None, original_im=original_pic.im)
            image_deal_info = {}
            x1, y1, x2, y2 = cut_image.getbbox()
            image_deal_info["鞋子原始位置"] = (x1, y1, x2, y2)
            o_w, o_h = cut_image.size
            image_deal_info["鞋子原始抠图后大小"] = (o_w, o_h)
            # 扩边处理
            _w, _h = x2 - x1, y2 - y1
            out_px = 0.06
            _w, _h = int(out_px * _w), int(out_px * _h)
            n_x1, n_y1, n_x2, n_y2 = x1 - _w, y1 - _h, x2 + _w, y2 + _h
            if n_x1 < 0:
                n_x1 = 0
            if n_y1 < 0:
                n_y1 = 0
            if n_x2 > o_w:
                n_x2 = o_w
            if n_y2 > o_h:
                n_y2 = o_h

            image_deal_info["抠图扩边后位置"] = (n_x1, n_y1, n_x2, n_y2)
            cut_image = original_pic.im.crop(image_deal_info["抠图扩边后位置"])
            image_deal_info["抠图扩边后图片大小"] = cut_image.size
            image_deal_info["原始图片大小"] = (original_pic.x, original_pic.y)

            # 使用pixian进行抠图
            second_cut_image, _ = self.direct_matting_image(image=cut_image)
            if not second_cut_image:
                return False, {"data": _}

            if second_cut_image.size != image_deal_info["抠图扩边后图片大小"]:
                print("图片尺寸还原")
                second_cut_image = second_cut_image.resize(image_deal_info["抠图扩边后图片大小"])

            # 创建空白图片并粘贴回去
            _img_im = Image.new(mode="RGBA", size=image_deal_info["原始图片大小"], color=(0, 0, 0, 0))
            _img_im.paste(second_cut_image, box=(image_deal_info["抠图扩边后位置"][0], image_deal_info["抠图扩边后位置"][1]))
        else:
            _img_im = self.direct_matting_image(image=original_pic.im)
            pass

        if out_file_path:
            _img_im.save(out_file_path)
        return True, {}

    def download_picture(self, url, out_path):
        response = requests.get(url)
        pic = response.content
        with open(out_path, 'wb') as f:
            f.write(pic)


if __name__ == '__main__':
    r = RemoveBgPiXian()
    path = r"C:\Users\gymmc\Desktop\白底部分品类45度\测试\eva_keai_1.png"
    out_path = "{}._no_bg-out.png".format(path)
    r.get_image_cut(path, out_file_path=out_path)