Commit abd38c59 authored by Yuxin Wu's avatar Yuxin Wu

imgaug

parent 6f9f4cd9
...@@ -324,10 +324,10 @@ class LocallyShuffleData(ProxyDataFlow, RNGDataFlow): ...@@ -324,10 +324,10 @@ class LocallyShuffleData(ProxyDataFlow, RNGDataFlow):
self.rng.shuffle(self.q) self.rng.shuffle(self.q)
for _ in range(self.q.maxlen): for _ in range(self.q.maxlen):
yield self.q.popleft() yield self.q.popleft()
self.q.append(next(self.ds_itr))
cnt += 1 cnt += 1
if cnt == self.size(): if cnt == self.size():
return return
self.q.append(next(self.ds_itr))
def SelectComponent(ds, idxs): def SelectComponent(ds, idxs):
......
...@@ -90,7 +90,7 @@ class ILSVRCMeta(object): ...@@ -90,7 +90,7 @@ class ILSVRCMeta(object):
obj = caffepb.BlobProto() obj = caffepb.BlobProto()
mean_file = os.path.join(self.dir, 'imagenet_mean.binaryproto') mean_file = os.path.join(self.dir, 'imagenet_mean.binaryproto')
with open(mean_file) as f: with open(mean_file, 'rb') as f:
obj.ParseFromString(f.read()) obj.ParseFromString(f.read())
arr = np.array(obj.data).reshape((3, 256, 256)).astype('float32') arr = np.array(obj.data).reshape((3, 256, 256)).astype('float32')
arr = np.transpose(arr, [1,2,0]) arr = np.transpose(arr, [1,2,0])
...@@ -105,6 +105,7 @@ class ILSVRC12(DataFlow): ...@@ -105,6 +105,7 @@ class ILSVRC12(DataFlow):
""" """
assert name in ['train', 'test', 'val'] assert name in ['train', 'test', 'val']
self.full_dir = os.path.join(dir, name) self.full_dir = os.path.join(dir, name)
assert os.path.isdir(self.full_dir), self.full_dir
self.shuffle = shuffle self.shuffle = shuffle
self.meta = ILSVRCMeta(meta_dir) self.meta = ILSVRCMeta(meta_dir)
self.imglist = self.meta.get_image_list(name) self.imglist = self.meta.get_image_list(name)
......
...@@ -6,7 +6,11 @@ ...@@ -6,7 +6,11 @@
from .base import ImageAugmentor from .base import ImageAugmentor
__all__ = ['RandomChooseAug', 'MapImage'] __all__ = ['RandomChooseAug', 'MapImage', 'Identity']
class Identity(ImageAugmentor):
def _augment(self, img, _):
return img
class RandomChooseAug(ImageAugmentor): class RandomChooseAug(ImageAugmentor):
def __init__(self, aug_lists): def __init__(self, aug_lists):
......
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: noise.py
# Author: Yuxin Wu <ppwwyyxxc@gmail.com>
from .base import ImageAugmentor
import numpy as np
import cv2
__all__ = ['JpegNoise', 'GaussianNoise']
class JpegNoise(ImageAugmentor):
def __init__(self, quality_range=(40, 100)):
self._init(locals())
def _get_augment_params(self, img):
return self._rand_range(*self.quality_range)
def _augment(self, img, q):
return cv2.imdecode(cv2.imencode('.jpg', img,
[cv2.IMWRITE_JPEG_QUALITY, q])[1], 1)
class GaussianNoise(ImageAugmentor):
def __init__(self, scale=10, clip=True):
self._init(locals())
def _get_augment_params(self, img):
return self.rng.randn(img.shape)
def _augment(self, img, noise):
ret = img + noise
if self.clip:
ret = np.clip(ret, 0, 255)
return ret
...@@ -6,7 +6,7 @@ from .base import ImageAugmentor ...@@ -6,7 +6,7 @@ from .base import ImageAugmentor
import numpy as np import numpy as np
import cv2 import cv2
__all__ = ['Flip', 'Resize', 'RandomResize', 'JpegNoise'] __all__ = ['Flip', 'Resize', 'RandomResize']
class Flip(ImageAugmentor): class Flip(ImageAugmentor):
""" """
...@@ -58,7 +58,7 @@ class Resize(ImageAugmentor): ...@@ -58,7 +58,7 @@ class Resize(ImageAugmentor):
class RandomResize(ImageAugmentor): class RandomResize(ImageAugmentor):
""" randomly rescale w and h of the image""" """ randomly rescale w and h of the image"""
def __init__(self, xrange, yrange, minimum=None, aspect_ratio_thres=0.2): def __init__(self, xrange, yrange, minimum=(0,0), aspect_ratio_thres=0.15):
""" """
:param xrange: (min, max) scaling ratio :param xrange: (min, max) scaling ratio
:param yrange: (min, max) scaling ratio :param yrange: (min, max) scaling ratio
...@@ -82,14 +82,3 @@ class RandomResize(ImageAugmentor): ...@@ -82,14 +82,3 @@ class RandomResize(ImageAugmentor):
def _augment(self, img, dsize): def _augment(self, img, dsize):
return cv2.resize(img, dsize, interpolation=cv2.INTER_CUBIC) return cv2.resize(img, dsize, interpolation=cv2.INTER_CUBIC)
class JpegNoise(ImageAugmentor):
def __init__(self, quality_range=(40, 100)):
self._init(locals())
def _get_augment_params(self, img):
return self._rand_range(*self.quality_range)
def _augment(self, img, q):
return cv2.imdecode(cv2.imencode('.jpg', img,
[cv2.IMWRITE_JPEG_QUALITY, q])[1], 1)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment