diff --git a/paddle_hub/data/augmentation.py b/paddle_hub/data/augmentation.py index 9d0531501ca43921438ee5b2fb58ac0ad2396d1b..1fef8cf94fc9736c044bb8a1521b169adafefe46 100644 --- a/paddle_hub/data/augmentation.py +++ b/paddle_hub/data/augmentation.py @@ -11,3 +11,205 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from PIL import Image, ImageEnhance +from paddle_hub.tools import utils +import numpy as np + + +def _check_range_0_1(value): + value = value if value <= 1 else 1 + value = value if value >= 0 else 0 + return value + + +def _check_bound(low, high): + low = _check_range_0_1(low) + high = _check_range_0_1(high) + high = high if high >= low else low + return low, high + + +def _check_img(img): + if isinstance(img, str): + utils.check_path(img) + img = Image.open(img) + return img + + +def _check_img_and_size(img, width, height): + img = _check_img(img) + img_width, img_height = img.size + height = height if img_height > height else img_height + height = img_height if height <= 0 else height + width = width if img_width > width else img_width + width = img_width if width <= 0 else width + return img, width, height + + +def image_crop_from_position(img, width, height, w_start, h_start): + img, width, height = _check_img_and_size(img, width, height) + w_end = w_start + width + h_end = h_start + height + return img.crop((w_start, h_start, w_end, h_end)) + + +def image_crop_from_TL(img, width, height): + w_start = h_start = 0 + return image_crop_from_position(img, width, height, w_start, h_start) + + +def image_crop_from_TR(img, width, height): + img, width, height = _check_img_and_size(img, width, height) + w_start = img.size[0] - width + h_start = 0 + return image_crop_from_position(img, width, height, w_start, h_start) + + +def image_crop_from_BL(img, width, height): + img, width, height = _check_img_and_size(img, width, height) + w_start = 0 + h_start = img.size[1] - height + return image_crop_from_position(img, width, height, w_start, h_start) + + +def image_crop_from_BR(img, width, height): + img, width, height = _check_img_and_size(img, width, height) + w_start = img.size[0] - width + h_start = img.size[1] - height + return image_crop_from_position(img, width, height, w_start, h_start) + + +def image_crop_from_centor(img, width, height): + img = _check_img(img) + w_start = (img.size[0] - width) / 2 + h_start = (img.size[1] - height) / 2 + return image_crop_from_position(img, width, height, w_start, h_start) + + +def image_crop_random(img, width=0, height=0): + img = _check_img(img) + width = width if width else np.random.randint( + int(img.size[0] / 10), img.size[0]) + height = height if height else np.random.randint( + int(img.size[1] / 10), img.size[1]) + w_start = np.random.randint(0, img.size[0] - width) + h_start = np.random.randint(0, img.size[1] - height) + return image_crop_from_position(img, width, height, w_start, h_start) + + +def image_resize(img, width, height, interpolation_method=Image.LANCZOS): + img = _check_img(img) + return img.resize((width, height), interpolation_method) + + +def image_resize_random(img, + width=0, + height=0, + interpolation_method=Image.LANCZOS): + img = _check_img(img) + width = width if width else np.random.randint( + int(img.size[0] / 10), img.size[0]) + height = height if height else np.random.randint( + int(img.size[1] / 10), img.size[1]) + return image_resize(img, width, height, interpolation_method) + + +def image_rotate(img, angle, expand=False): + img = _check_img(img) + return img.rotate(angle, expand=expand) + + +def image_rotate_random(img, low=0, high=360, expand=False): + angle = np.random.randint(low, high) + return image_rotate(img, angle, expand) + + +def image_brightness_adjust(img, delta): + delta = _check_range_0_1(delta) + img = _check_img(img) + return ImageEnhance.Brightness(img).enhance(delta) + + +def image_brightness_adjust_random(img, low=0, high=1): + low, high = _check_bound(low, high) + delta = np.random.uniform(low, high) + return image_brightness_adjust(img, delta) + + +def image_contrast_adjust(img, delta): + delta = _check_range_0_1(delta) + img = _check_img(img) + return ImageEnhance.Contrast(img).enhance(delta) + + +def image_contrast_adjust_random(img, low=0, high=1): + low, high = _check_bound(low, high) + delta = np.random.uniform(low, high) + return image_contrast_adjust(img, delta) + + +def image_saturation_adjust(img, delta): + delta = _check_range_0_1(delta) + img = _check_img(img) + return ImageEnhance.Color(img).enhance(delta) + + +def image_saturation_adjust_random(img, low=0, high=1): + low, high = _check_bound(low, high) + delta = np.random.uniform(low, high) + return image_saturation_adjust(img, delta) + + +def image_flip_top_bottom(img): + img = _check_img(img) + return img.transpose(Image.FLIP_TOP_BOTTOM) + + +def image_flip_left_right(img): + img = _check_img(img) + return img.transpose(Image.FLIP_LEFT_RIGHT) + + +def image_flip_random(img): + img = _check_img(img) + flag = np.random.randint(0, 1) + if flag: + return image_flip_top_bottom(img) + else: + return image_flip_left_right(img) + + +def image_random_process(img, + enable_resize=True, + enable_crop=True, + enable_rotate=True, + enable_brightness_adjust=True, + enable_contrast_adjust=True, + enable_saturation_adjust=True, + enable_flip=True): + operator_list = [] + if enable_resize: + operator_list.append(image_resize_random) + if enable_crop: + operator_list.append(image_crop_random) + if enable_rotate: + operator_list.append(image_rotate_random) + if enable_brightness_adjust: + operator_list.append(image_brightness_adjust_random) + if enable_contrast_adjust: + operator_list.append(image_contrast_adjust_random) + if enable_saturation_adjust: + operator_list.append(image_saturation_adjust_random) + if enable_flip: + operator_list.append(image_flip_random) + + if not operator_list: + return img + + random_op_index = np.random.randint(0, len(operator_list) - 1) + random_op = operator_list[random_op_index] + return random_op(img)