python 图像数据扩充代码(RandomColor / RandomRotate / PicinPic 等)

  • picinpic(将logo随机加在背景上)
# coding:utf-8
from __future__ import print_function, absolute_import
import threading
import os
import gc
import cv2
import random
import numpy as np
#imutilsimport
import sys

from PIL import ImageFile
from PIL import Image, ImageFont, ImageDraw, ImageFilter, ImageEnhance
from imgaug import augmenters as iaa
from sklearn.preprocessing import normalize as sknormalize
import multiprocessing
from math import pi
from utils import pil_pad_image, shrink
from copy import deepcopy





class BaseAugmenter(object):
    def __init__(self):
        pass 
    def augment_image(self,img):
        pass
    def augment_images(self,imgs):
        return [self.augment_image(img) for img in imgs ]


class RandomResize(BaseAugmenter):
    '''
    当scale是一个float
    当scale长度为2,则为最大值和最小值,在其中随机生成。
    当scale长度大于2,则为指定的几个值,在其中随机选

    当fix_pos = True时,按照示例图中的样子: 直接pad时,缩放后的图在黑边的最下边中央位置
    '''
    def __init__(self, scale=(0.5,0.9),keep_size=False,pad=True,fix_pos=True):
        self.scale = scale
        self.keep_size = False 
        self.pad = pad
        self.fix_pos = fix_pos
    def augment_image(self,img):
        _w, _h = img.width, img.height
        if isinstance(self.scale, float):
            h = int(self.scale*_h)
            w = int(self.scale*_w)
            new_img = img.resize((w,h))
        elif not self.keep_size:
            if len(self.scale) == 2:
                h = int(random.uniform(self.scale[0],self.scale[1])*_h)
                w = int(random.uniform(self.scale[0],self.scale[1])*_w)
            elif len(self.scale) > 2:
                h = int(random.choice(self.scale)*_h)
                w = int(random.choice(self.scale)*_w)
            new_img = img.resize((w,h))
        else:
            if len(self.scale) == 2:
                ratio = random.uniform(self.scale[0], self.scale[1])
            elif len(self.scale) >2:
                ratio = random.choice(self.scale)
            h,w = int(ratio*_h), int(ratio*_w)
            new_img = img.resize((w,h))
        
        if self.pad:
            bg = Image.new(new_img.mode, (_w, _h))
            if not self.fix_pos:
                x = random.randint(0, bg.width - new_img.width)
                y = random.randint(0, bg.height - new_img.height) 
                bg.paste(new_img, (x,y))
            else:
                x = int((bg.width-new_img.width)/2)
                y = bg.height - new_img.height
                bg.paste(new_img, (x,y))
            new_img = bg
        return new_img


class PicInPic(BaseAugmenter): 
    def __init__(self,bgimg_list_fname,keep_size,thresh1,thresh2,angle_r,angle_s):
        self.keep_size = keep_size
        self.thresh1 = thresh1
        self.thresh2 = thresh2
        self.angle = angle_r
        self.shear = iaa.Affine(shear=angle_s)
        self.bg_list = open(bgimg_list_fname).read().splitlines()
#        self.affine = iaa.Affine(
#            scale={"x": (0.8, 1.2), "y": (0.8, 1.2)}, # scale images to 80-120% of their size, individually per axis
#            translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)}, # translate by -20 to +20 percent (per axis)
#            rotate=(-45, 45), # rotate by -45 to +45 degrees
#            shear=(-30, 30), # shear by -16 to +16 degrees
#            order=[0, 1], # use nearest neighbour or bilinear interpolation (fast)
#            cval=(0, 255), # if mode is constant, use a cval between 0 and 255
#            mode=ia.ALL # use any of scikit-image's warping modes (see 2nd image from the top for examples)
#        )
 
    def affine(self, img):
        if img.mode == "RGB":
            new_img = img.convert("RGBA")
        elif img.mode == "P":
            new_img = img.convert("RGBA")
        else:
            # print(">>>>>>>>>>>>>", img.mode)
            new_img = img.convert("RGBA")
        new_img = pil_pad_image(new_img,3,4)
        if random.random() > self.thresh1:
            angle = random.uniform(self.angle[0], self.angle[1])
            new_img = new_img.rotate(angle)
            if random.random() > self.thresh2:
                new_img = Image.fromarray(self.shear.augment_image(np.array(new_img)))
        elif random.random() > self.thresh1:
            new_img = Image.fromarray(self.shear.augment_image(np.array(new_img)))
        new_img = shrink(new_img)
        return new_img

    def augment_image(self,img):
        bg_name = random.choice(self.bg_list)
        bg = Image.open(bg_name)

        r_wh = bg.width / bg.height

        while r_wh < 0.4 or r_wh > 2.5:
            print(bg_name, " r_wh yichang")
            # bg.show()
            bg_name = random.choice(self.bg_list)
            bg = Image.open(bg_name)
            r_wh = bg.width / bg.height

        bg = bg.convert("RGBA")
        # bg.show()
        # img.show()
        if self.keep_size:
            bg = bg.resize((img.width,img.height))
        new_img = self.affine(img)

        # new_img.show()
        ratio = min(float(bg.height)/new_img.height, float(bg.width)/new_img.width)
        # r = random.uniform(0.65*ratio, 0.9*ratio)
        r = random.uniform(0.7*ratio, 0.95*ratio)
        w,h = int(r*new_img.width), int(r*new_img.height)
        new_img = new_img.resize((w,h))
        r,g,b,a = new_img.split()

        mode_x = random.randint(0, 1)
        mode_y = random.randint(0, 1)

        if mode_x == 0:
            x = random.randint(0, int(0.25*(bg.width - new_img.width)))
        elif mode_x == 1:
            x = random.randint(int(0.75*(bg.width - new_img.width)), bg.width - new_img.width)


        if mode_y == 0:
            y = random.randint(0, int(0.25*(bg.height - new_img.height)))
        elif mode_y == 1:
            y = random.randint(int(0.75*(bg.height - new_img.height)), bg.height - new_img.height)




        # x = random.randint(0, bg.width - new_img.width)
        # y = random.randint(0, bg.height - new_img.height)
        # bgr = bg.paste(new_img,(x,y),mask=r)
        # bgr.show()
        # bgg = bg.paste(new_img,(x,y),mask=g)
        # bgg.show()
        # bgb = bg.paste(new_img,(x,y),mask=b)
        # bgb.show()
        # bga = bg.paste(new_img,(x,y),mask=a)
        # bga.show()
        bg.paste(new_img,(x,y),mask=a)
        # bg.show()
        # bg.show()
        return bg, bg_name

class Watermark(BaseAugmenter):
    def __init__(self, logolist_fname):
        self.logo_files = open(logolist_fname).read().splitlines()
    def augment_image(self, img):
        logo = Image.open(random.choice(self.logo_files))
        logo = logo.convert("RGBA")
        assert logo.mode == "RGBA"
        logo = pil_pad_image(logo,3,4)
        logo = logo.rotate(random.uniform(-180,180))
        logo = shrink(logo)

        ratio = min(float(img.height)/logo.height,  float(img.width)/logo.width)
        ratio = random.uniform(0.4,0.8)*ratio
        h,w = int(ratio*logo.height), int(ratio*logo.width)
        logo = logo.resize((w,h))
        x = random.randint(0, img.width - logo.width)
        y = random.randint(0, img.height - logo.height)

        r,g,b,a = logo.split()
        img.paste(logo, (x,y), mask=a)
        return img

class Mosaic(BaseAugmenter):
    def augment_image(self, img):
        if isinstance(img, Image.Image):
            img = np.array(img)
        img_height,img_width = img.shape[:2]
        random_x = random.uniform(0.3, 0.6)
        random_y = random.uniform(0.3, 0.6)
        masaic_width = int(img_width*random_x)
        masaic_height = int(img_height*random_y)
        paste_coordx = random.randint(0, img_width - masaic_width)
        paste_coordy = random.randint(0, img_height - masaic_height)
        size = random.randint(10,25)
        width_bin = int(masaic_width/float(size))
        height_bin = int(masaic_height/float(size))
        for x in range(width_bin):
            for y in range(height_bin):
                img[paste_coordy+y*size:paste_coordy+(y+1)*size,paste_coordx+x*size:paste_coordx+(x+1)*size,:] = img[paste_coordy+y*size,paste_coordx+x*size,:]
        return Image.fromarray(img)

class SaltNoise(BaseAugmenter):
    def __init__(self, p):
        self.func = iaa.SaltAndPepper(p=p)
    def augment_image(self,img):
        if isinstance(img, Image.Image):
            img = np.array(img)
        img = self.func.augment_image(img)
        return Image.fromarray(img)

class WhiteNoise(BaseAugmenter):
    def __init__(self,loc,scale,per_channel):
        self.func = iaa.AdditiveGaussianNoise(loc=loc,scale=scale,per_channel=per_channel)
    def augment_image(self,img):
        if isinstance(img, Image.Image):
            img = np.array(img)
        img = self.func.augment_image(img)
        return Image.fromarray(img)

class ColorToGray(BaseAugmenter):
    def __init__(self):
        if random.random()>0.1:
            self.func = iaa.Sequential([
            iaa.Grayscale(alpha=(0.1,1.0)),
            iaa.ContrastNormalization((0.5, 2.0), per_channel=0.5),
            iaa.Add(value=(-45,45))
            ],
            random_order=False)
        else:
            self.func = iaa.Sequential([
            iaa.Grayscale(alpha=(0.5,1.0)),
            ],
            random_order=False)
    def augment_image(self,img):
        if isinstance(img, Image.Image):
            img = np.array(img)
        img = self.func.augment_image(img)
        return Image.fromarray(img)



def all_augment(img, k):
    results = []
    # for func, name in zip([pic_in_pic,mosaic,color_to_gray,watermark], ["picInPic","mosaic","gray","watermark"]):
    for func, name in zip([pic_in_pic,mosaic,color_to_gray,watermark], ["picInPic"]):
        for i in range(k):
            #
            # if image_src.size[0] < 200 and image_src.size[1] < 200:
            #     print(img_name, " is too small")
            #     continue



            # res_img = random_resize.augment_image(deepcopy(img))
            res_img, bg_name = func.augment_image(deepcopy(img))
            results.append((name + '_' + str(i), res_img, bg_name))
    # for func, name in zip([white_noise, salt_noise], ["white_noise","salt_noise"]):
    #     for i in range(5):
    #         res_img = func.augment_image(deepcopy(img))
    #         results.append((name + '_' + str(i), res_img))
    # for i in range(10):
    #     res_img = random_resize.augment_image(deepcopy(img))
    #     if i % 2 == 0:
    #         res_img = res_img.transpose(Image.FLIP_LEFT_RIGHT)
    #     else:
    #         res_img = res_img.transpose(Image.FLIP_TOP_BOTTOM)
    #     results.append(("random_resize_" + str(i),res_img))

    return results


# for bg_fix in ["fengmian", "neiye"]:
for bg_fix in ["fengmian", "neiye"]:
    BG_LIST_FILE = "bg_v6_" + bg_fix + "_aug.txt"


    # BG_LIST_FILE = "bg_v6_fengmian_aug.txt"
    # LOGO_LIST_FILE = "logo_v6_kuang_dianzi.txt"
    # LOGO_LIST_FILE = "logo_v6_kuang_shengcheng_text.txt"
    # LOGO_LIST_FILE = "logo_v6_kuang_shouxie_rectango.txt"

    # for LOGO_fix in ["dianzi", "shouxie_square", "logo_v6_kuang_shouxie_rectango_xi", "shengcheng_text", "yinshua" ]:
    for LOGO_fix in ["shengcheng_text"]:


        LOGO_LIST_FILE = "logo_v6_kuang_" + LOGO_fix + ".txt"
        # LOGO_LIST_FILE = "logo_v6_kuang_yinshua.txt"

        # IMG_LIST_FILE = "/home/netease/code/homologous_image_retrieval/data/imgs.txt"

        if bg_fix == "fengmian":
            if LOGO_fix == "dianzi": k = 26
            elif LOGO_fix == "shouxie_square": k = 59
            elif LOGO_fix == "shouxie_rectango": k = 48
            elif LOGO_fix == "shengcheng_text": k = 22
            elif LOGO_fix == "yinshua": k = 20
            elif LOGO_fix == "shouxie_rectango_xi": k = 30
            elif LOGO_fix == "shouxie_rectango_cu": k = 30

        if bg_fix == "neiye":
            if LOGO_fix == "dianzi": k = 12
            elif LOGO_fix == "shouxie_square": k = 36
            elif LOGO_fix == "shouxie_rectango": k = 30
            elif LOGO_fix == "shengcheng_text": k = 8
            elif LOGO_fix == "yinshua": k = 10
            elif LOGO_fix == "shouxie_rectango_xi": k = 30
            elif LOGO_fix == "shouxie_rectango_cu": k = 30


        k = 15




        IMG_LIST_FILE = LOGO_LIST_FILE

        salt_noise = SaltNoise(p=(0.05, 0.25))
        white_noise = WhiteNoise(loc=0, scale=(0.05 * 255, 0.20 * 255), per_channel=0.5)
        pic_in_pic = PicInPic(bgimg_list_fname=BG_LIST_FILE, keep_size=False, thresh1=0.5, thresh2=0.5, angle_r=(-70, 70),
                              angle_s=(-0, 0))
        random_resize = RandomResize(scale=[0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.9])
        mosaic = Mosaic()
        color_to_gray = ColorToGray()

        watermark = Watermark(LOGO_LIST_FILE)

        i=0
        curPath = os.path.abspath('.')

        for imgname in open(IMG_LIST_FILE).read().splitlines():
            img = Image.open(imgname)
            print(imgname)




            results = all_augment(img, k)
            # save_path = "5output_goodlogo_v2/"+str(i)+'/'
            save_path = "5output_goodlogo_v10_new/" + bg_fix + "_kuang_" + LOGO_fix + "_07to095_shear0_" + str(k) + "/"
            save_path = os.path.join(curPath, save_path)

            # print(os.path.exists(save_path))

            if not os.path.exists(save_path):
                # os.mkdir(save_path)
                os.makedirs(save_path)
            for name, res_img, bg_name in results:
                # res_name = save_path + imgname.split('/')[-1][:-4] + '_' + name + '.png'
                res_name = save_path + imgname.split('/')[-1][:-4] + '_' + bg_name.split('/')[-1][:-4] + '.png'
                res_img.save(res_name)
            i+=1







需要的utils:

import gc
import cv2
import random
import numpy as np
import sys

from PIL import ImageFile
from PIL import Image, ImageFont, ImageDraw, ImageFilter, ImageEnhance
from imgaug import augmenters as iaa
from sklearn.preprocessing import normalize as sknormalize
import multiprocessing
from math import pi

def pil_pad_image(img, rate,c=4):
    _h,_w = img.height, img.width
    w = rate*_w
    h = rate*_h
    mode = {3:"RGB", 4:"RGBA"}[c]
    new_img = Image.new(mode, (w,h))
    y = int((h - _h)/2)
    x = int((w - _w)/2)
    new_img.paste(img, ((x,y)))
    return new_img

def shrink(img):
    if not isinstance(img, np.ndarray):
        img = np.array(img)
    mask = img>0
    coords = np.argwhere(mask)
    if len(img.shape)==3:
        y0,x0,_ = coords.min(axis = 0)
        y1,x1,_ = coords.max(axis = 0)+1
        new_img = img[y0:y1,x0:x1,:]
    else:
        y0,x0 = coords.min(axis = 0)
        y1,x1 = coords.max(axis = 0)+1
        new_img = img[y0:y1,x0:x1]
    return Image.fromarray(np.uint8(new_img))
  • randomColor \ randomRotate \ randomCrop \ randomGaussian
import matplotlib.pyplot as plt
from PIL import Image
import random
import os
# image_path = "/home/xsr-ai/datasets/butterfly.jpg"
from PIL import Image, ImageEnhance, ImageOps, ImageFile
import numpy as np
import random
import threading, os, time
import logging

logger = logging.getLogger(__name__)
ImageFile.LOAD_TRUNCATED_IMAGES = True


class DataAugmentation:
    """
    包含数据增强的八种方式
    """


    def __init__(self):
        pass

    @staticmethod
    def openImage(image):
        return Image.open(image, mode="r")

    @staticmethod
    def randomRotation(image, mode=Image.BICUBIC):
        """
         对图像进行随机任意角度(0~360度)旋转
        :param mode 邻近插值,双线性插值,双三次B样条插值(default)
        :param image PIL的图像image
        :return: 旋转转之后的图像
        """
        random_angle = np.random.randint(1, 5)
        return image.rotate(random_angle, mode)

    @staticmethod
    def randomCrop(image):
        """
        对图像随意剪切,考虑到图像大小范围(68,68),使用一个一个大于(36*36)的窗口进行截图
        :param image: PIL的图像image
        :return: 剪切之后的图像

        """
        image_width = image.size[0]
        image_height = image.size[1]

        crop_width = int(image_width*0.81649658092)
        crop_height = int(image_height*0.81649658092)

        # crop_win_size = np.random.randint(40, 68)
        random_region = (
            (image_width - crop_width) >> 1, (image_height - crop_height) >> 1, (image_width + crop_width) >> 1,
            (image_height + crop_height) >> 1)
        return image.crop(random_region)

    @staticmethod
    def randomColor(image):
        """
        对图像进行颜色抖动
        :param image: PIL的图像image
        :return: 有颜色色差的图像image
        """
        # random_factor = np.random.randint(8, 12) / 10.  # 随机因子
        # color_image = ImageEnhance.Color(image).enhance(random_factor)  # 调整图像的饱和度
        # random_factor = np.random.randint(8, 12) / 10.  # 随机因子
        # brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor)  # 调整图像的亮度
        # random_factor = np.random.randint(8, 15) / 10.  # 随机因1子
        # contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)  # 调整图像对比度
        # random_factor = np.random.randint(8, 12) / 10.  # 随机因子

        random_factor = np.random.randint(5, 15) / 10.  # 随机因子
        color_image = ImageEnhance.Color(image).enhance(random_factor)  # 调整图像的饱和度
        random_factor = np.random.randint(8, 13) / 10.  # 随机因子
        brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor)  # 调整图像的亮度
        random_factor = np.random.randint(7, 17) / 10.  # 随机因1子
        contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)  # 调整图像对比度
        random_factor = np.random.randint(5, 16) / 10.  # 随机因子


        return ImageEnhance.Sharpness(contrast_image).enhance(random_factor)  # 调整图像锐度

    @staticmethod
    def randomGaussian(image, mean=0.2, sigma=0.3):
        """
         对图像进行高斯噪声处理
        :param image:
        :return:
        """

        def gaussianNoisy(im, mean=0.2, sigma=0.3):
            """
            对图像做高斯噪音处理
            :param im: 单通道图像
            :param mean: 偏移量
            :param sigma: 标准差
            :return:
            """
            for _i in range(len(im)):
                im[_i] += random.gauss(mean, sigma)
            return im

        # 将图像转化成数组
        img = np.asarray(image)
        img.flags.writeable = True  # 将数组改为读写模式
        width, height = img.shape[:2]
        img_r = gaussianNoisy(img[:, :, 0].flatten(), mean, sigma)
        img_g = gaussianNoisy(img[:, :, 1].flatten(), mean, sigma)
        img_b = gaussianNoisy(img[:, :, 2].flatten(), mean, sigma)
        img[:, :, 0] = img_r.reshape([width, height])
        img[:, :, 1] = img_g.reshape([width, height])
        img[:, :, 2] = img_b.reshape([width, height])
        return Image.fromarray(np.uint8(img))

    @staticmethod
    def saveImage(image, path):
        image.save(path)


def imageOps(func_name, image, des_path, file_name, times=1):
    funcMap = {"randomRotation": DataAugmentation.randomRotation,
               "randomCrop": DataAugmentation.randomCrop,
               "randomColor": DataAugmentation.randomColor,
               "randomGaussian": DataAugmentation.randomGaussian
               }
    if funcMap.get(func_name) is None:
        logger.error("%s is not exist", func_name)
        return -1

    for _i in range(0, times, 1):
        new_image = funcMap[func_name](image)
        DataAugmentation.saveImage(new_image, os.path.join(des_path, func_name + str(_i) + file_name))


# opsList = {"randomRotation", "randomCrop", "randomColor", "randomGaussian"}
opsList = {"randomColor", "randomRotation"}


def threadOPS(image, crop_img_name):

    threadImage = [0] * 2
    _index = 0

    for ops_name in opsList:
        threadImage[_index] = threading.Thread(target=imageOps,
                                               args=(ops_name, image, save_path, crop_img_name))
        threadImage[_index].start()
        _index += 1
        time.sleep(0.2)

curPath = os.path.abspath('.')



def random_crop(image, crop_shape, padding=None):
    oshape = image.size

    if padding:
        oshape_pad = (oshape[0] + 2 * padding, oshape[1] + 2 * padding)
        img_pad = Image.new("RGB", (oshape_pad[0], oshape_pad[1]))
        img_pad.paste(image, (padding, padding))

        nh = random.randint(0, oshape_pad[0] - crop_shape[0])
        nw = random.randint(0, oshape_pad[1] - crop_shape[1])
        image_crop = img_pad.crop((nh, nw, nh + crop_shape[0], nw + crop_shape[1]))

        return image_crop
    else:
        print("WARNING!!! nothing to do!!!")
        return image


if __name__ == "__main__":


    # for bgname in ["fm", "ny"]:
    for bgname in ["fm"]:

        # for prename in ['yinshua_da', 'dianzi', 'yinshua_xiao', "shouxie"]:
        for prename in ["yinshua_da_shai"]:
            #
            # file_dir = "/media/crx/4f5a4954-cbbd-46a9-9c48-aeb156476154/netease/file/1program/3ad/0banzheng/pachongpic/google/0nologo_v6_all_" + bgname + "/"
            file_dir = "/media/crx/4f5a4954-cbbd-46a9-9c48-aeb156476154/netease/file/1program/3ad/0banzheng/17173_inference/17173_ny_dianzi_692_2of3/"
            # file_dir = os.path.join(curPath, bgname + "_" + prename)
            # save_path = "/media/crx/4f5a4954-cbbd-46a9-9c48-aeb156476154/netease/file/1program/3ad/0banzheng/pachongpic/google/0nologo_v7_all_" + bgname + "_aug/"
            save_path = "/media/crx/4f5a4954-cbbd-46a9-9c48-aeb156476154/netease/file/1program/3ad/0banzheng/17173_inference/17173_ny_dianzi_692_2of3_aug4/"
            # save_path = os.path.join(curPath, bgname + "_" + prename + "_aug")

            if not os.path.exists(save_path):
                # os.mkdir(save_path)
                os.makedirs(save_path)


            for img_name in os.listdir(file_dir):
                img_path = os.path.join(curPath, file_dir, img_name)

                image_src = Image.open(img_path)
                # crop_width = image_src.size[0] - 24
                # crop_height = image_src.size[1] - 24



                # crop_width = int(image_src.size[0]*0.81649658092)
                # crop_height = int(image_src.size[1]*0.81649658092)
                if image_src.size[0] < 200 and image_src.size[1] < 200:
                    print(img_name, " is too small")
                    continue
                else:

                    for i in range(2):

                        r = random.random()

                        if r <= 0.5:
                            image_src = image_src.transpose(Image.FLIP_LEFT_RIGHT)
                            if r <= 0.25:
                                image_src = image_src.transpose(Image.FLIP_TOP_BOTTOM)
                        if r > 0.5:
                            image_src = image_src.transpose(Image.ROTATE_90)
                            if r < 0.75:
                                image_src = image_src.transpose(Image.ROTATE_180)
                        # image_src = im.transpose(image_src.ROTATE_270)

                        crop_width = random.randint(int(image_src.size[0]*0.81649658092), int(image_src.size[0]*1))
                        crop_height = random.randint(int(image_src.size[1]*0.81649658092), int(image_src.size[1]*1))


                        image_dst_crop = random_crop(image_src, [crop_width, crop_height], padding=10)

                        # im = img_name.split('/')[-1][:-4]

                        # crop_img_path = os.path.join(curPath, save_path, "crop" + str(i+1) + img_name)
                        crop_img_path = os.path.join(curPath, save_path, img_name.split('/')[-1][:-5] + "_crop" + str(i+1) + ".jpg")
                        crop_img_name = img_name.split('/')[-1][:-4] + "_crop" + str(i+1) + ".jpg"
                        # image_dst_crop.save(crop_img_path)

                        # print("here!")
                        # threadOPS(image_dst_crop, crop_img_name)
                        threadOPS(image_dst_crop, crop_img_name)


  • image shrink(去掉透明边缘)
import cv2
import numpy as np
import sys
import os
# import sets
import multiprocessing

def shrink(cv_img):
    mask = cv_img>0
    coords = np.argwhere(mask)
    if len(cv_img.shape)==3:
        y0,x0,_ = coords.min(axis = 0)
        y1,x1,_ = coords.max(axis = 0)+1
        new_img = cv_img[y0:y1,x0:x1,:]
    else:
        y0,x0 = coords.min(axis = 0)
        y1,x1 = coords.max(axis = 0)+1
        new_img = cv_img[y0:y1,x0:x1]
    return new_img


# logo_path = '2logo_demo/'
logo_path = '/media/crx/4f5a4954-cbbd-46a9-9c48-aeb156476154/netease/file/1program/3ad/0banzheng/pachongpic/google/2logo_v6集合/kuang_dianzi/'

#logo_name_list = [os.path.join(logo_path, logo_name) for logo_name in os.listdir(logo_path) if os.path.isfile(os.path.join(logo_path, logo_name))]
logo_name_list = []
for root,dirs,files in os.walk(logo_path):
    for f in files:
        path = os.path.join(root,f)
        if os.path.isfile(path):
            logo_name_list.append(path)
'''
logo_list = {}
for logo_name in logo_name_list:
    for logo in LOGO_LIST:
        if logo in logo_name:
            if logo not in logo_list:
                logo_list[logo] = [logo_name]
            else:
                logo_list[logo].append(logo_name)
            break
'''
savedir = logo_path+'output/'
print(savedir)

def start(begin,end):
    for i in range(begin,end):
        image_name = logo_name_list[i]
        print(image_name)
        try:
            img = cv2.imread(image_name,-1)
            new_img = shrink(img)
            cv2.imwrite(image_name,new_img)
        except:
            print('error!')
            os.remove(image_name)
            continue

multiprocess=2
# if multiprocess:
# 	num_workers = 4
# 	part = int(len(logo_name_list)/num_workers)
# 	id_min = [i*part for i in range(num_workers)]
# 	id_max = [(i+1)*part for i in range(num_workers)]
# 	pool = multiprocessing.Pool(processes=num_workers)
# 	for i in range(num_workers):
# 		pool.apply_async(start,args=(id_min[i], id_max[i]))
# 	pool.close()
# 	pool.join()
# else:
# 	start(0,len(logo_name_list))

start(0,len(logo_name_list))

猜你喜欢

转载自blog.csdn.net/s000da/article/details/90242740
今日推荐