mirrors/fast-reid
1
0
Fork 0
mirror of https://github.com/JDAI-CV/fast-reid.git synced 2025-06-03 14:50:47 +08:00
Code Issues Projects Releases Wiki Activity

fix randompatch (#470)

Browse Source
This commit is contained in:
Xingyu Liao 2021-04-21 17:08:16 +08:00 committed by GitHub
parent 0c8e3d9805
commit 46b0681313
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
fastreid/data/transforms/transforms.py
View File

@ -11,7 +11,7 @@ import random
from collections import deque from collections import deque
import numpy as np import numpy as np
from PIL import Image import torch
from .functional import to_tensor, augmentations from .functional import to_tensor, augmentations
@ -55,8 +55,7 @@ class RandomPatch(object):
""" """
def __init__(self, prob_happen=0.5, pool_capacity=50000, min_sample_size=100, def __init__(self, prob_happen=0.5, pool_capacity=50000, min_sample_size=100,
patch_min_area=0.01, patch_max_area=0.5, patch_min_ratio=0.1, patch_min_area=0.01, patch_max_area=0.5, patch_min_ratio=0.1, prob_flip_leftright=0.5,
prob_rotate=0.5, prob_flip_leftright=0.5,
): ):
self.prob_happen = prob_happen self.prob_happen = prob_happen
@ -64,7 +63,6 @@ class RandomPatch(object):
self.patch_max_area = patch_max_area self.patch_max_area = patch_max_area
self.patch_min_ratio = patch_min_ratio self.patch_min_ratio = patch_min_ratio
self.prob_rotate = prob_rotate
self.prob_flip_leftright = prob_flip_leftright self.prob_flip_leftright = prob_flip_leftright
self.patchpool = deque(maxlen=pool_capacity) self.patchpool = deque(maxlen=pool_capacity)
@ -83,23 +81,18 @@ class RandomPatch(object):
def transform_patch(self, patch): def transform_patch(self, patch):
if random.uniform(0, 1) > self.prob_flip_leftright: if random.uniform(0, 1) > self.prob_flip_leftright:
patch = patch.transpose(Image.FLIP_LEFT_RIGHT) patch = torch.flip(patch, dims=[2])
if random.uniform(0, 1) > self.prob_rotate:
patch = patch.rotate(random.randint(-10, 10))
return patch return patch
def __call__(self, img): def __call__(self, img):
if isinstance(img, np.ndarray): _, H, W = img.size() # original image size
img = Image.fromarray(img.astype(np.uint8))
W, H = img.size # original image size
# collect new patch # collect new patch
w, h = self.generate_wh(W, H) w, h = self.generate_wh(W, H)
if w is not None and h is not None: if w is not None and h is not None:
x1 = random.randint(0, W - w) x1 = random.randint(0, W - w)
y1 = random.randint(0, H - h) y1 = random.randint(0, H - h)
new_patch = img.crop((x1, y1, x1 + w, y1 + h)) new_patch = img[..., y1:y1 + h, x1:x1 + w]
self.patchpool.append(new_patch) self.patchpool.append(new_patch)
if len(self.patchpool) < self.min_sample_size: if len(self.patchpool) < self.min_sample_size:
@ -110,11 +103,11 @@ class RandomPatch(object):
# paste a randomly selected patch on a random position # paste a randomly selected patch on a random position
patch = random.sample(self.patchpool, 1)[0] patch = random.sample(self.patchpool, 1)[0]
patchW, patchH = patch.size _, patchH, patchW = patch.size()
x1 = random.randint(0, W - patchW) x1 = random.randint(0, W - patchW)
y1 = random.randint(0, H - patchH) y1 = random.randint(0, H - patchH)
patch = self.transform_patch(patch) patch = self.transform_patch(patch)
img.paste(patch, (x1, y1)) img[..., y1:y1 + patchH, x1:x1 + patchW] = patch
return img return img