train/src/masks_integrated.py

import math
import random
import logging
from enum import Enum

import cv2
import numpy as np
import random

LOGGER = logging.getLogger(__name__)

class LinearRamp:
    def __init__(self, start_value=0, end_value=1, start_iter=-1, end_iter=0):
        self.start_value = start_value
        self.end_value = end_value
        self.start_iter = start_iter
        self.end_iter = end_iter

    def __call__(self, i):
        if i < self.start_iter:
            return self.start_value
        if i >= self.end_iter:
            return self.end_value
        part = (i - self.start_iter) / (self.end_iter - self.start_iter)
        return self.start_value * (1 - part) + self.end_value * part

class DrawMethod(Enum):
    LINE = 'line'
    CIRCLE = 'circle'
    SQUARE = 'square'

def make_random_irregular_mask(shape, max_angle=4, max_len=60, max_width=20, min_times=0, max_times=10,
                               draw_method=DrawMethod.LINE):
    """生成不规则mask - 基于角度和长度的线条"""
    draw_method = DrawMethod(draw_method)

    height, width = shape
    mask = np.zeros((height, width), np.float32)
    times = np.random.randint(min_times, max_times + 1)
    for i in range(times):
        start_x = np.random.randint(width)
        start_y = np.random.randint(height)
        for j in range(1 + np.random.randint(5)):
            angle = 0.01 + np.random.randint(max_angle)
            if i % 2 == 0:
                angle = 2 * 3.1415926 - angle
            length = 10 + np.random.randint(max_len)
            brush_w = 5 + np.random.randint(max_width)
            end_x = np.clip((start_x + length * np.sin(angle)).astype(np.int32), 0, width)
            end_y = np.clip((start_y + length * np.cos(angle)).astype(np.int32), 0, height)
            if draw_method == DrawMethod.LINE:
                cv2.line(mask, (start_x, start_y), (end_x, end_y), 1.0, brush_w)
            elif draw_method == DrawMethod.CIRCLE:
                cv2.circle(mask, (start_x, start_y), radius=brush_w, color=1., thickness=-1)
            elif draw_method == DrawMethod.SQUARE:
                radius = brush_w // 2
                mask[start_y - radius:start_y + radius, start_x - radius:start_x + radius] = 1
            start_x, start_y = end_x, end_y
    return mask[None, ...]


def make_random_rectangle_mask(shape, margin=10, bbox_min_size=30, bbox_max_size=100, min_times=0, max_times=3):
    """生成随机矩形mask"""
    height, width = shape
    mask = np.zeros((height, width), np.float32)
    bbox_max_size = min(bbox_max_size, height - margin * 2, width - margin * 2)
    times = np.random.randint(min_times, max_times + 1)
    for i in range(times):
        box_width = np.random.randint(bbox_min_size, bbox_max_size)
        box_height = np.random.randint(bbox_min_size, bbox_max_size)
        start_x = np.random.randint(margin, width - margin - box_width + 1)
        start_y = np.random.randint(margin, height - margin - box_height + 1)
        mask[start_y:start_y + box_height, start_x:start_x + box_width] = 1
    return mask[None, ...]


def make_random_superres_mask(shape, min_step=2, max_step=4, min_width=1, max_width=3):
    """生成超分辨率风格的规则网格mask"""
    height, width = shape
    mask = np.zeros((height, width), np.float32)
    step_x = np.random.randint(min_step, max_step + 1)
    width_x = np.random.randint(min_width, min(step_x, max_width + 1))
    offset_x = np.random.randint(0, step_x)

    step_y = np.random.randint(min_step, max_step + 1)
    width_y = np.random.randint(min_width, min(step_y, max_width + 1))
    offset_y = np.random.randint(0, step_y)

    for dy in range(width_y):
        mask[offset_y + dy::step_y] = 1
    for dx in range(width_x):
        mask[:, offset_x + dx::step_x] = 1
    return mask[None, ...]


def make_brush_stroke_mask(shape, num_strokes_range=(1, 5), stroke_width_range=(5, 30), 
                          max_offset=50, num_points_range=(5, 15)):
    """生成笔刷描边样式的mask - 基于随机偏移的连续线条"""
    num_strokes = random.randint(*num_strokes_range)
    height, width = shape
    mask = np.zeros((height, width), dtype=np.float32)
    
    for _ in range(num_strokes):
        # 随机起点
        start_x = random.randint(0, width)
        start_y = random.randint(0, height)
        
        # 随机描边参数
        num_points = random.randint(*num_points_range)
        stroke_width = random.randint(*stroke_width_range)
        
        points = [(start_x, start_y)]
        
        # 生成连续的点
        for i in range(num_points):
            prev_x, prev_y = points[-1]
            # 添加随机偏移
            dx = random.randint(-max_offset, max_offset)
            dy = random.randint(-max_offset, max_offset)
            new_x = max(0, min(width, prev_x + dx))
            new_y = max(0, min(height, prev_y + dy))
            points.append((new_x, new_y))
        
        # 绘制描边
        for i in range(len(points) - 1):
            cv2.line(mask, points[i], points[i+1], 1.0, stroke_width)
    
    return mask[None, ...]


class RandomIrregularMaskGenerator:
    """不规则mask生成器"""
    def __init__(self, max_angle=4, max_len=60, max_width=20, min_times=0, max_times=10, ramp_kwargs=None,
                 draw_method=DrawMethod.LINE):
        self.max_angle = max_angle
        self.max_len = max_len
        self.max_width = max_width
        self.min_times = min_times
        self.max_times = max_times
        self.draw_method = draw_method
        self.ramp = LinearRamp(**ramp_kwargs) if ramp_kwargs is not None else None

    def __call__(self, img, iter_i=None, raw_image=None):
        coef = self.ramp(iter_i) if (self.ramp is not None) and (iter_i is not None) else 1
        cur_max_len = int(max(1, self.max_len * coef))
        cur_max_width = int(max(1, self.max_width * coef))
        cur_max_times = int(self.min_times + 1 + (self.max_times - self.min_times) * coef)
        return make_random_irregular_mask(img.shape[1:], max_angle=self.max_angle, max_len=cur_max_len,
                                          max_width=cur_max_width, min_times=self.min_times, max_times=cur_max_times,
                                          draw_method=self.draw_method)


class RandomRectangleMaskGenerator:
    """矩形mask生成器"""
    def __init__(self, margin=10, bbox_min_size=30, bbox_max_size=100, min_times=0, max_times=3, ramp_kwargs=None):
        self.margin = margin
        self.bbox_min_size = bbox_min_size
        self.bbox_max_size = bbox_max_size
        self.min_times = min_times
        self.max_times = max_times
        self.ramp = LinearRamp(**ramp_kwargs) if ramp_kwargs is not None else None

    def __call__(self, img, iter_i=None, raw_image=None):
        coef = self.ramp(iter_i) if (self.ramp is not None) and (iter_i is not None) else 1
        cur_bbox_max_size = int(self.bbox_min_size + 1 + (self.bbox_max_size - self.bbox_min_size) * coef)
        cur_max_times = int(self.min_times + (self.max_times - self.min_times) * coef)
        return make_random_rectangle_mask(img.shape[1:], margin=self.margin, bbox_min_size=self.bbox_min_size,
                                          bbox_max_size=cur_bbox_max_size, min_times=self.min_times,
                                          max_times=cur_max_times)


class RandomSuperresMaskGenerator:
    """超分辨率mask生成器"""
    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def __call__(self, img, iter_i=None):
        return make_random_superres_mask(img.shape[1:], **self.kwargs)


class BrushStrokeMaskGenerator:
    """笔刷描边mask生成器"""
    def __init__(self, num_strokes_range=(1, 5), stroke_width_range=(5, 30), 
                 max_offset=50, num_points_range=(5, 15), ramp_kwargs=None):
        self.num_strokes_range = num_strokes_range
        self.stroke_width_range = stroke_width_range
        self.max_offset = max_offset
        self.num_points_range = num_points_range
        self.ramp = LinearRamp(**ramp_kwargs) if ramp_kwargs is not None else None

    def __call__(self, img, iter_i=None, raw_image=None):
        coef = self.ramp(iter_i) if (self.ramp is not None) and (iter_i is not None) else 1
        cur_num_strokes = int(max(1, self.num_strokes_range[1] * coef))
        cur_stroke_width_range = (
            int(max(1, self.stroke_width_range[0] * coef)),
            int(max(1, self.stroke_width_range[1] * coef))
        )
        return make_brush_stroke_mask(
            img.shape[1:], 
            num_strokes_range=(cur_num_strokes, cur_num_strokes),
            stroke_width_range=cur_stroke_width_range,
            max_offset=self.max_offset,
            num_points_range=self.num_points_range
        )


class DumbAreaMaskGenerator:
    """简单区域mask生成器"""
    min_ratio = 0.1
    max_ratio = 0.35
    default_ratio = 0.225

    def __init__(self, is_training):
        #Parameters:
        #    is_training(bool): If true - random rectangular mask, if false - central square mask
        self.is_training = is_training

    def _random_vector(self, dimension):
        if self.is_training:
            lower_limit = math.sqrt(self.min_ratio)
            upper_limit = math.sqrt(self.max_ratio)
            mask_side = round((random.random() * (upper_limit - lower_limit) + lower_limit) * dimension)
            u = random.randint(0, dimension-mask_side-1)
            v = u+mask_side 
        else:
            margin = (math.sqrt(self.default_ratio) / 2) * dimension
            u = round(dimension/2 - margin)
            v = round(dimension/2 + margin)
        return u, v

    def __call__(self, img, iter_i=None, raw_image=None):
        c, height, width = img.shape
        mask = np.zeros((height, width), np.float32)
        x1, x2 = self._random_vector(width)
        y1, y2 = self._random_vector(height)
        mask[x1:x2, y1:y2] = 1
        return mask[None, ...]


class IntegratedMaskGenerator:
    """集成的mask生成器 - 支持多种mask类型混合"""
    def __init__(self, irregular_proba=1/4, irregular_kwargs=None,
                 box_proba=1/4, box_kwargs=None,
                 segm_proba=1/4, segm_kwargs=None,
                 brush_stroke_proba=1/4, brush_stroke_kwargs=None,
                 superres_proba=0, superres_kwargs=None,
                 squares_proba=0, squares_kwargs=None,
                 invert_proba=0):
        self.probas = []
        self.gens = []

        if irregular_proba > 0:
            self.probas.append(irregular_proba)
            if irregular_kwargs is None:
                irregular_kwargs = {}
            else:
                irregular_kwargs = dict(irregular_kwargs)
            irregular_kwargs['draw_method'] = DrawMethod.LINE
            self.gens.append(RandomIrregularMaskGenerator(**irregular_kwargs))

        if box_proba > 0:
            self.probas.append(box_proba)
            if box_kwargs is None:
                box_kwargs = {}
            self.gens.append(RandomRectangleMaskGenerator(**box_kwargs))

        if brush_stroke_proba > 0:
            self.probas.append(brush_stroke_proba)
            if brush_stroke_kwargs is None:
                brush_stroke_kwargs = {}
            self.gens.append(BrushStrokeMaskGenerator(**brush_stroke_kwargs))

        if superres_proba > 0:
            self.probas.append(superres_proba)
            if superres_kwargs is None:
                superres_kwargs = {}
            self.gens.append(RandomSuperresMaskGenerator(**superres_kwargs))

        if squares_proba > 0:
            self.probas.append(squares_proba)
            if squares_kwargs is None:
                squares_kwargs = {}
            else:
                squares_kwargs = dict(squares_kwargs)
            squares_kwargs['draw_method'] = DrawMethod.SQUARE
            self.gens.append(RandomIrregularMaskGenerator(**squares_kwargs))

        self.probas = np.array(self.probas, dtype='float32')
        self.probas /= self.probas.sum()
        self.invert_proba = invert_proba

    def __call__(self, img, iter_i=None, raw_image=None):
        kind = np.random.choice(len(self.probas), p=self.probas)
        gen = self.gens[kind]
        result = gen(img, iter_i=iter_i, raw_image=raw_image)
        if self.invert_proba > 0 and random.random() < self.invert_proba:
            result = 1 - result
        return result


def get_mask_generator(kind, kwargs):
    """获取mask生成器的工厂函数"""
    if kind is None:
        kind = "integrated"
    if kwargs is None:
        kwargs = {}

    if kind == "integrated":
        cl = IntegratedMaskGenerator
    elif kind == "irregular":
        cl = RandomIrregularMaskGenerator
    elif kind == "rectangle":
        cl = RandomRectangleMaskGenerator
    elif kind == "brush_stroke":
        cl = BrushStrokeMaskGenerator
    elif kind == "superres":
        cl = RandomSuperresMaskGenerator
    elif kind == "dumb":
        cl = DumbAreaMaskGenerator
    else:
        raise NotImplementedError(f"No such generator kind = {kind}")
    return cl(**kwargs)