From 4601778d258636a52d6f52fb404602cf1c2adda1 Mon Sep 17 00:00:00 2001
From: pyhqos7bg <wangwei_jmu@163.com>
Date: Thu, 30 May 2024 15:25:09 +0800
Subject: [PATCH] ADD file via upload

---
 augment.py | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 augment.py

diff --git a/augment.py b/augment.py
new file mode 100644
index 0000000..de95ac2
--- /dev/null
+++ b/augment.py
@@ -0,0 +1,79 @@
+import os
+import shutil
+import cv2
+import numpy as np
+ROOT_PATH="./train_data/"
+SAVE_PATH='./aug_data/'
+
+unified_width=250
+unified_height=250
+#递归删除该目录下的所有文件夹及子文件夹
+if os.path.exists(SAVE_PATH):
+    shutil.rmtree(SAVE_PATH)
+#创建多层文件夹，用于存储增强图像数据
+def make_directory(read_path=ROOT_PATH,save_path=SAVE_PATH):
+    subfolders = os.listdir(read_path)
+    for subfolder in subfolders:
+        sub_save_path = save_path+subfolder+'/'
+        os.makedirs(sub_save_path)
+
+#将文件夹下的所有图片统一某个尺寸
+def resize_images(read_path=ROOT_PATH,save_path=SAVE_PATH,width = unified_width,height = unified_height):
+    subfolders = os.listdir(read_path)
+    for subfolder in subfolders:
+        file_names = os.listdir(read_path+subfolder+'/')
+        for filename in file_names:
+            path_read_filename = read_path+subfolder+"/"+filename
+            img = cv2.imread(path_read_filename)
+            img = cv2.resize(img,(width,height),interpolation=cv2.INTER_CUBIC)
+            path_save_file = save_path+subfolder+"/"+"resized_"+filename
+            cv2.imwrite(path_save_file,img)
+
+#直方图均衡化处理
+def equalize_hist_images(read_path = ROOT_PATH,save_path=SAVE_PATH):
+    subfolders = os.listdir(read_path)
+    for subfolder in subfolders:
+        filenames = os.listdir(read_path+subfolder+"/")
+        for filename in filenames:
+            file_read_path = read_path+subfolder+"/"+filename
+            img = cv2.imread(file_read_path)
+            b,g,r = cv2.split(img)
+            b = cv2.equalizeHist(b)
+            g = cv2.equalizeHist(g)
+            r = cv2.equalizeHist(r)
+            img = cv2.merge([b,g,r])
+            file_save_path = save_path+subfolder+"/"+"aug_equ_"+filename
+            cv2.imwrite(file_save_path,img)
+
+#旋转图片，并扩充保存到指定文件夹下
+def rotate_images(read_path = ROOT_PATH,save_path=SAVE_PATH,num_rotation=3):
+    rotation_angle = np.random.randint(-30,30)
+    subfolders= os.listdir(read_path)
+    for subfolder in subfolders:
+        filenames = os.listdir(read_path+subfolder+"/")
+        for filename in filenames:
+            img = cv2.imread(read_path+subfolder+"/"+filename)
+            rows,cols,channel = img.shape
+            M = cv2.getRotationMatrix2D(((cols-1)/2.0,(rows-1)/2.0),rotation_angle,1)
+            img_rotate = cv2.warpAffine(img,M,(cols,rows))
+            rotate_write_path = save_path+subfolder+"/"+"aug_rotate_"+filename
+            cv2.imwrite(rotate_write_path,img_rotate)
+
+#进行模糊处理
+def blur_images(read_path=ROOT_PATH,save_path=SAVE_PATH):
+    blur_size = np.random.randint(0,4)*2+1
+    subfolders = os.listdir(read_path)
+    for subfolder in subfolders:
+        filenames = os.listdir(read_path+subfolder+"/")
+        for filename in filenames:
+            img = cv2.imread(read_path+subfolder+"/"+filename)
+            img_blur = cv2.GaussianBlur(img,(blur_size,blur_size),0) #高斯模糊必须是整数
+            blur_file_path = save_path+subfolder+"/"+"Gaussian_"+filename
+            cv2.imwrite(blur_file_path,img_blur)
+
+if __name__ == '__main__':
+    make_directory(ROOT_PATH,SAVE_PATH)
+    resize_images()
+    rotate_images()
+    equalize_hist_images()
+    blur_images()
\ No newline at end of file