CheckCardBoxAndDefectServer/app/utils/card_inference/data_augmentation.py

import numpy as np
import math
import cv2


class LetterBox(object):
    def __init__(self, size={'width': 640, 'height': 640}, auto=False, stride=32, *args, **kwargs):
        # 需要调整的额size
        self.size = size
        self.h = size["height"]
        self.w = size["width"]
        self.auto = auto  # pass max size integer, automatically solve for short side using stride
        self.stride = stride  # used with auto

    def __call__(self, im_lb):
        imgList = im_lb['imgList']
        lb = im_lb['lb']
        if lb is not None:
            assert imgList[0].shape[:2] == lb.shape[:2]

        ans_imgList = self.handle_imgList(imgList)

        # 处理label
        # 处理label
        if lb is not None:
            ans_lb = self.handle_single_label(lb)
        else:
            ans_lb = None

        returnObj = dict(imgList=ans_imgList, lb=ans_lb)
        return returnObj

    def handle_imgList(self, imgList):
        # 处理图片
        ans_imgList = []
        for per_img in imgList:
            ans_img = self.handle_single_img(per_img)
            ans_imgList.append(ans_img)
        return ans_imgList

    def get_offset(self, originImgSize={'width': 4096, 'height': 7000}):

        # _240429_1543_
        # [特别注意]：ResizeBeforeLetterbox中重写了这个逻辑

        originH = originImgSize['height']
        originW = originImgSize['width']

        dstH = self.h
        dstW = self.w

        def fry_resize_realParams(originH, originW, dstH, dstW):
            r = min(dstH / originH, dstW / originW)  # ratio of new/old
            resize_h, resize_w = int(round(originH * r)), int(round(originW * r))  # resized image
            total_pad_h = int(dstH - resize_h)
            total_pad_w = int(dstW - resize_w)
            assert total_pad_h >= 0, "total_pad_h 必须大于等于0"
            assert total_pad_w >= 0, "total_pad_w 必须大于等于0"

            assert total_pad_h == 0 or total_pad_w == 0, "total_pad_h 和 total_pad_w中必须有一个为0"

            pad_left = int(total_pad_w // 2)
            pad_right = total_pad_w - pad_left
            pad_top = int(total_pad_h // 2)
            pad_bottom = total_pad_h - pad_top

            before_letterbox_dict = {}
            before_letterbox_dict['ratio'] = r
            before_letterbox_dict['resize_h'] = resize_h
            before_letterbox_dict['resize_w'] = resize_w
            before_letterbox_dict['total_pad_h'] = total_pad_h
            before_letterbox_dict['total_pad_w'] = total_pad_w
            before_letterbox_dict['pad_left'] = pad_left
            before_letterbox_dict['pad_right'] = pad_right
            before_letterbox_dict['pad_top'] = pad_top
            before_letterbox_dict['pad_bottom'] = pad_bottom

            return before_letterbox_dict

        before_letterbox_dict = fry_resize_realParams(originH, originW, dstH, dstW)

        rect_dict = {}
        rect_dict['x'] = before_letterbox_dict['pad_left']
        rect_dict['y'] = before_letterbox_dict['pad_top']
        rect_dict['width'] = before_letterbox_dict['resize_w']
        rect_dict['height'] = before_letterbox_dict['resize_h']
        rect_dict['ratio'] = before_letterbox_dict['ratio']
        return rect_dict

    def handle_single_img(self, im):
        assert len(im.shape) == 3, "im 必须是3维的"
        assert (im.shape[2] == 1) or (im.shape[2] == 3), "im 的通道数必须是一个通道或者三个通道"

        imh, imw = im.shape[:2]

        r = min(self.h / imh, self.w / imw)  # ratio of new/old
        h, w = round(imh * r), round(imw * r)  # resized image
        hs, ws = (math.ceil(x / self.stride) * self.stride for x in (h, w)) if self.auto else self.h, self.w
        top, left = round((hs - h) / 2 - 0.1), round((ws - w) / 2 - 0.1)

        # 这里弄成0没有关系，因为均值是0方差是1
        # 还是都弄成114吧
        if im.shape[2] == 3:
            im_out = np.full((self.h, self.w, 3), 114, dtype=im.dtype)
        elif im.shape[2] == 1:
            im_out = np.full((self.h, self.w, 1), 114, dtype=im.dtype)
        else:
            raise ValueError("图片的通道数异常")

        if im.shape[2] == 1:
            gray_image_hw1 = im
            gray_image_hw = np.squeeze(gray_image_hw1, axis=-1)
            singleImg = gray_image_hw
        else:
            singleImg = im

        originImg_resized = cv2.resize(singleImg, (w, h), interpolation=cv2.INTER_LINEAR)

        if len(originImg_resized.shape) == 2:
            newSingleImg2D = originImg_resized
            newSingleImg3D = np.expand_dims(newSingleImg2D, axis=-1)
            newSingleImg = newSingleImg3D
        else:
            newSingleImg = originImg_resized

        im_out[top:top + h, left:left + w] = newSingleImg

        return im_out

    def handle_single_label(self, im):
        assert len(im.shape) == 2, "label 必须是2维的"

        imh, imw = im.shape[:2]

        r = min(self.h / imh, self.w / imw)  # ratio of new/old
        h, w = round(imh * r), round(imw * r)  # resized image
        hs, ws = (math.ceil(x / self.stride) * self.stride for x in (h, w)) if self.auto else self.h, self.w
        top, left = round((hs - h) / 2 - 0.1), round((ws - w) / 2 - 0.1)

        # label 直接弄成255不参与计算
        im_out = np.full((self.h, self.w), 0, dtype=im.dtype)
        im_out[top:top + h, left:left + w] = cv2.resize(im, (w, h), interpolation=cv2.INTER_NEAREST)

        return im_out