mirror of https://github.com/JDAI-CV/fast-reid.git
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
# encoding: utf-8
|
|
"""
|
|
@author: liaoxingyu
|
|
@contact: liaoxingyu2@jd.com
|
|
"""
|
|
|
|
from fastai.vision import *
|
|
from fastai.vision.image import *
|
|
|
|
|
|
def _random_erasing(x, probability=0.5, sl=0.02, sh=0.4, r1=0.3,
|
|
mean=(np.array(imagenet_stats[1]) + 1) * imagenet_stats[0]):
|
|
if random.uniform(0, 1) > probability:
|
|
return x
|
|
|
|
for attempt in range(100):
|
|
area = x.size()[1] * x.size()[2]
|
|
|
|
target_area = random.uniform(sl, sh) * area
|
|
aspect_ratio = random.uniform(r1, 1 / r1)
|
|
|
|
h = int(round(math.sqrt(target_area * aspect_ratio)))
|
|
w = int(round(math.sqrt(target_area / aspect_ratio)))
|
|
|
|
if w < x.size()[2] and h < x.size()[1]:
|
|
x1 = random.randint(0, x.size()[1] - h)
|
|
y1 = random.randint(0, x.size()[2] - w)
|
|
if x.size()[0] == 3:
|
|
x[0, x1:x1 + h, y1:y1 + w] = mean[0]
|
|
x[1, x1:x1 + h, y1:y1 + w] = mean[1]
|
|
x[2, x1:x1 + h, y1:y1 + w] = mean[2]
|
|
else:
|
|
x[0, x1:x1 + h, y1:y1 + w] = mean[0]
|
|
return x
|
|
|
|
|
|
RandomErasing = TfmPixel(_random_erasing)
|