Skip to content
Snippets Groups Projects
Commit 6304b647 authored by Korabelnikov Aleks's avatar Korabelnikov Aleks Committed by Kai Chen
Browse files

generalization of Expand transfrom probability (#1633)

Allows to use not only 0.5 probability.
parent 4b101ece
No related branches found
No related tags found
No related merge requests found
......@@ -509,13 +509,15 @@ class Expand(object):
mean (tuple): mean value of dataset.
to_rgb (bool): if need to convert the order of mean to align with RGB.
ratio_range (tuple): range of expand ratio.
prob (float): probability of applying this transformation
"""
def __init__(self,
mean=(0, 0, 0),
to_rgb=True,
ratio_range=(1, 4),
seg_ignore_label=None):
seg_ignore_label=None,
prob=0.5):
self.to_rgb = to_rgb
self.ratio_range = ratio_range
if to_rgb:
......@@ -524,9 +526,10 @@ class Expand(object):
self.mean = mean
self.min_ratio, self.max_ratio = ratio_range
self.seg_ignore_label = seg_ignore_label
self.prob = prob
def __call__(self, results):
if random.randint(2):
if random.uniform(0, 1) > self.prob:
return results
img, boxes = [results[k] for k in ('img', 'gt_bboxes')]
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment