一蝇庭、做隨機亮度、對比度、飽和度修改秉继,使用tensorflow API
核心部分是 aug_op 函數(shù)祈噪,這可是菜鳥的心血啊尚辑!
#coding:utf-8
import tensorflow as tf
import cv2
import random
import sys
import os
import shutil
#os.environ["CUDA_VISIBLE_DEVICES"] = ""
def random_normal(img, mean=0.0, stddev=2.0, limit=6):
img = tf.cast(img, dtype=tf.int32)# This type cast is very important.
noise = tf.random_normal(tf.shape(img), mean, stddev)
noise = tf.cast(noise, dtype=tf.int32)# This type cast is very important. For negative value.
up_limit = tf.where(tf.less_equal(noise, limit), tf.ones_like(noise, dtype=tf.uint8), tf.zeros_like(noise, dtype=tf.uint8))
down_limit = tf.where(tf.greater_equal(noise, 0-limit), tf.ones_like(noise, dtype=tf.uint8), tf.zeros_like(noise, dtype=tf.uint8))
up_limit = tf.cast(up_limit, dtype=tf.int32)
down_limit = tf.cast(down_limit, dtype=tf.int32)
noise = (up_limit*down_limit)*noise
new_img = tf.add(img, noise)
new_img = tf.cast(img, dtype=tf.uint8)
return new_img
def aug_op(img):
new_img = img
cnt = 0
a = tf.constant([0,1],dtype=tf.int32)
b = tf.random_shuffle(a)
cnt = tf.constant([0], dtype=tf.int32)
new_img, cnt = tf.cond(b[0]>0, lambda: (tf.image.random_brightness(new_img, 0.3), cnt+1), lambda: (new_img,cnt))
new_img, cnt = tf.cond(b[0]>0, lambda: (tf.image.random_contrast(new_img, 0.5, 2), cnt+1), lambda: (new_img,cnt))
new_img, cnt = tf.cond(b[0]>0, lambda: (tf.image.random_saturation(new_img, 0.3, 2), cnt+1) , lambda: (new_img,cnt))
new_img, cnt = tf.cond(b[0]>0, lambda: (random_normal(new_img), cnt+1), lambda: (new_img,cnt+1))
return new_img, cnt
if __name__ == "__main__":
src_dir = sys.argv[1]
src_landmark_dir = sys.argv[2]
dst_dir = sys.argv[3]
dst_landmark_dir = sys.argv[4]
loop_times = 5
if not os.path.exists(src_dir):
print("{} not exists,exit.".format(src_dir))
if not os.path.exists(src_landmark_dir):
print("{} not exists,exit.".format(src_landmark_dir))
if not os.path.exists(dst_dir):
os.makedirs(dst_dir)
if not os.path.exists(dst_landmark_dir):
os.makedirs(dst_landmark_dir)
sess = tf.Session()
input_image = tf.placeholder(tf.uint8, shape=[None, None, 3], name="input_image")
op, cnt_op = aug_op(input_image)
src_names = os.listdir(src_dir)
cnt = 0
total_cnt = len(src_names)
new_cnt = 0
for name in src_names:
img = cv2.imread(os.path.join(src_dir, name))
if img is None:
print("Image:{} read none, skip".format(os.path.join(src_dir, name)))
continue
pre, ext = os.path.splitext(name)
pts_name = pre+".pts"
for i in range(loop_times):
new_name = pre+"_"+str(i)+ext
new_pts_name = pre+"_"+str(i)+".pts"
new_img, isaug= sess.run([op,cnt_op], feed_dict={input_image:img})
print(isaug)
if isaug == 0:
continue
cv2.imwrite(os.path.join(dst_dir, new_name), new_img)
shutil.copy(os.path.join(src_landmark_dir, pts_name), os.path.join(dst_landmark_dir, new_pts_name))
new_cnt = new_cnt + 1
cnt = cnt + 1
print("progress: {}/{}, new image:{}".format(cnt, total_cnt, new_cnt))
#python img_aug.py /home/hongjie.li/FaceAlign/dataset-landmark/20181018-1000-camera /home/hongjie.li/FaceAlign/dataset-landmark/landmark_all /home/hongjie.li/FaceAlign/dataset-landmark/20181018-1000-camera-aug /home/hongjie.li/FaceAlign/dataset-landmark/20181018-1000-camera-aug-landmark