Spaces:
Sleeping
Sleeping
File size: 5,681 Bytes
d4b77ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import random
import numpy as np
class RandomMask():
def __init__(self, videoLength, dataInfo):
self.videoLength = videoLength
self.imageHeight, self.imageWidth = dataInfo['image']['image_height'], \
dataInfo['image']['image_width']
self.maskHeight, self.maskWidth = dataInfo['mask']['mask_height'], \
dataInfo['mask']['mask_width']
try:
self.maxDeltaHeight, self.maxDeltaWidth = dataInfo['mask']['max_delta_height'], \
dataInfo['mask']['max_delta_width']
except KeyError:
self.maxDeltaHeight, self.maxDeltaWidth = 0, 0
try:
self.verticalMargin, self.horizontalMargin = dataInfo['mask']['vertical_margin'], \
dataInfo['mask']['horizontal_margin']
except KeyError:
self.verticalMargin, self.horizontalMargin = 0, 0
def __call__(self):
from .utils import random_bbox
from .utils import bbox2mask
masks = []
bbox = random_bbox(self.imageHeight, self.imageWidth, self.verticalMargin, self.horizontalMargin,
self.maskHeight, self.maskWidth)
if random.uniform(0, 1) > 0.5:
mask = bbox2mask(self.imageHeight, self.imageWidth, 0, 0, bbox)
for frame in range(self.videoLength):
masks.append(mask)
else:
for frame in range(self.videoLength):
delta_h, delta_w = random.randint(-3, 3), random.randint(-3, 3) # 每次向四个方向移动三个像素以内
bbox = list(bbox)
bbox[0] = min(max(self.verticalMargin, bbox[0] + delta_h), self.imageHeight - self.verticalMargin - bbox[2])
bbox[1] = min(max(self.horizontalMargin, bbox[1] + delta_w), self.imageWidth - self.horizontalMargin - bbox[3])
mask = bbox2mask(self.imageHeight, self.imageWidth, 0, 0, bbox)
masks.append(mask)
masks = np.stack(masks, axis=0)
if len(masks.shape) == 3:
masks = masks[:, :, :, np.newaxis]
assert len(masks.shape) == 4, 'Wrong mask dimension {}'.format(len(masks.shape))
return masks
class MidRandomMask():
### This mask is considered without random motion
def __init__(self, videoLength, dataInfo):
self.videoLength = videoLength
self.imageHeight, self.imageWidth = dataInfo['image']['image_height'], \
dataInfo['image']['image_width']
self.maskHeight, self.maskWidth = dataInfo['mask']['mask_height'], \
dataInfo['mask']['mask_width']
def __call__(self):
from .utils import mid_bbox_mask
mask = mid_bbox_mask(self.imageHeight, self.imageWidth, self.maskHeight, self.maskWidth)
masks = []
for _ in range(self.videoLength):
masks.append(mask)
return mask
class MatrixMask():
### This mask is considered without random motion
def __init__(self, videoLength, dataInfo):
self.videoLength = videoLength
self.imageHeight, self.imageWidth = dataInfo['image']['image_height'], \
dataInfo['image']['image_width']
self.maskHeight, self.maskWidth = dataInfo['mask']['mask_height'], \
dataInfo['mask']['mask_width']
try:
self.row, self.column = dataInfo['mask']['row'], \
dataInfo['mask']['column']
except KeyError:
self.row, self.column = 5, 4
def __call__(self):
from .utils import matrix2bbox
mask = matrix2bbox(self.imageHeight, self.imageWidth, self.maskHeight,
self.maskWidth, self.row, self.column)
masks = []
for video in range(self.videoLength):
masks.append(mask)
return mask
class FreeFormMask():
def __init__(self, videoLength, dataInfo):
self.videoLength = videoLength
self.imageHeight, self.imageWidth = dataInfo['image']['image_height'], \
dataInfo['image']['image_width']
self.maxVertex = dataInfo['mask']['max_vertex']
self.maxLength = dataInfo['mask']['max_length']
self.maxBrushWidth = dataInfo['mask']['max_brush_width']
self.maxAngle = dataInfo['mask']['max_angle']
def __call__(self):
from .utils import freeFormMask
mask = freeFormMask(self.imageHeight, self.imageWidth,
self.maxVertex, self.maxLength,
self.maxBrushWidth, self.maxAngle)
return mask
class StationaryMask():
def __init__(self, videoLength, dataInfo):
self.videoLength = videoLength
self.imageHeight, self.imageWidth = dataInfo['image']['image_height'], \
dataInfo['image']['image_width']
# self.maxPointNum = dataInfo['mask']['max_point_num']
# self.maxLength = dataInfo['mask']['max_length']
def __call__(self):
from .STTN_mask import create_random_shape_with_random_motion
masks = create_random_shape_with_random_motion(self.videoLength, 0.9, 1.1, 1, 10, self.imageHeight, self.imageWidth)
masks = np.stack(masks, axis=0)
if len(masks.shape) == 3:
masks = masks[:, :, :, np.newaxis]
assert len(masks.shape) == 4, 'Your masks with a wrong shape {}'.format(len(masks.shape))
return masks |