-
Notifications
You must be signed in to change notification settings - Fork 0
/
3_4_augment_only_rotate.py
executable file
·127 lines (104 loc) · 5.53 KB
/
3_4_augment_only_rotate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import os
import cv2
import albumentations as A
from tqdm import tqdm
import shutil
from multiprocessing import Pool
folderToAugmentJpgs = 'jpgs'
folderToAugmentTxts = 'txts'
augmentationExtension = 'only_rotate_random'
#define the transformations that we want to do
transform = A.Compose([
A.Rotate(limit=(-15,15),p=0.5)
], bbox_params=A.BboxParams(format='yolo'))
label='singlewood'
folderAfterAugmentationJpgs = folderToAugmentJpgs+'_'+augmentationExtension
folderAfterAugmentationTxts = folderToAugmentTxts+'_'+augmentationExtension
coresAmount = os.cpu_count()
#coresAmount = 1
database_base_path = '/home/jean-pierre/scratch/'
paths = [
database_base_path+'20201117_rotated/1pi2/labeled_20201117_1pi2_c3',
database_base_path+'20201117_rotated/2pi4/labeled_20201117_2pi4_c1',
database_base_path+'20201117_rotated/3SamsungGalaxyA5/labeled_20201117_3SamsungGalaxyA4_OC',
database_base_path+'20201117_rotated/4XiaomiRedmi4X/labeled_20201117_4XiaomiRedmi4X_OC',
database_base_path+'20201119_rotated/1XiaomiRedmi4X/labeled_20201119_1XiaomiRedmi4X_OC',
database_base_path+'20201119_rotated/2pi2/labeled_20201119_2pi2_c3',
database_base_path+'20201119_rotated/3pi4/labeled_20201119_3pi4_c3',
database_base_path+'20201119_rotated/4SamsungGalaxyA5_sampled/labeled_20201119_4SamsungGalaxyA_OC_sampled',
database_base_path+'20201126_rotated/3XiaomiRedmi2_sampled/labeled_20201126_3XiaomiRedmi2_OC2_sampled',
database_base_path+'20201126_rotated/4pi4_c4_sampled/labeled_20201126_4pi4_c4_sampled',
database_base_path+'20201126_rotated/4pi4_c5/labeled_20201126_4pi4_c5',
database_base_path+'20201127_rotated/2pi4/labeled_20201127_2pi4_c2',
database_base_path+'20201203_rotated/1pi4_sampled/labeled_20201203_1pi4_c20_sampled',
database_base_path+'20201203_rotated/3pi2/labeled_20201203_3pi2_c4',
database_base_path+'20201203_rotated/4SamsungGalaxyA5/labeled_20201203_4SamsungGalaxyA5_OC',
database_base_path+'unilyon_and_others/20071123_0756_Ain1/labeled_20071123_0756_Ain1',
database_base_path+'unilyon_and_others/20071123_0956_Ain2/labeled_20071123_0956_Ain2',
database_base_path+'unilyon_and_others/20191125_Allier1/labeled_20191125_Allier1',
database_base_path+'unilyon_and_others/20191223_Allier2/labeled_20191223_Allier2',
database_base_path+'unilyon_and_others/randomWoodImages/labeled_2022_randomWoodImages'
]
#create new txt files from bboxes (formatted: [(xcoor,ycoor,xwidth,yheight,'label'),(xcoor,ycoor,xwidth,yheight,'label'),(xcoor,ycoor,xwidth,yheight,'label')])
def saveNewYoloBboxes(newTxtFilename, newBboxesYolo):
#create text file
newTxtFile = open(newTxtFilename,"w")
#Loop through new bboxes
for newbbox in newBboxesYolo:
#write 6 digit floats per coordinate
newTxtFile.write("0 "+str("%.6f" % newbbox[0])+" "+str("%.6f" % newbbox[1])+" "+str("%.6f" % newbbox[2])+" "+str("%.6f" % newbbox[3])+"\n")
newTxtFile.close()
#loop through all paths
#for path in paths:
def augmentAllImagesInPath(path):
#create the directories to store the files in
try:
os.mkdir(os.path.join(path,folderAfterAugmentationJpgs))
except:
shutil.rmtree(os.path.join(path,folderAfterAugmentationJpgs))
os.mkdir(os.path.join(path,folderAfterAugmentationJpgs))
try:
os.mkdir(os.path.join(path,folderAfterAugmentationTxts))
except:
shutil.rmtree(os.path.join(path,folderAfterAugmentationTxts))
os.mkdir(os.path.join(path,folderAfterAugmentationTxts))
#loop through all files in the 'jpgs' folder
print(os.path.join(path,folderToAugmentJpgs))
for jpg in os.listdir(os.path.join(path,folderToAugmentJpgs)):
#print(jpg)
#create empty variable to store the labels in
yolo_coors = []
#determine the jpg and text files
if jpg[-4:] == '.jpg':
#print('jpg')
jpgFileName = jpg
txtFileName = jpg.replace('.jpg', '.txt')
jpgFileNameAugmented = jpgFileName.replace('.jpg', '_'+augmentationExtension+'.jpg')
txtFileNameAugmented = txtFileName.replace('.txt', '_'+augmentationExtension+'.txt')
jpgFilePath = os.path.join(path,folderToAugmentJpgs,jpgFileName)
txtFilePath = os.path.join(path,folderToAugmentTxts,txtFileName)
augmentedJpgPath = os.path.join(path,folderAfterAugmentationJpgs,jpgFileNameAugmented)
augmentedTxtPath = os.path.join(path,folderAfterAugmentationTxts,txtFileNameAugmented)
#read txt file
with open(txtFilePath) as f:
lines = f.readlines()
for line in lines:
coordinates = line.split(' ')
yolo_coors.append([float(coordinates[1]),float(coordinates[2]),float(coordinates[3]),float(coordinates[4].split('\n')[0]),label])
#read image
image = cv2.imread(jpgFilePath)
try:
#transform
transformed = transform(image=image, bboxes=yolo_coors)
transformed_image = transformed['image']
transformed_bboxes = transformed['bboxes']
cv2.imwrite(augmentedJpgPath,transformed_image)
saveNewYoloBboxes(augmentedTxtPath,transformed_bboxes)
except:
#print(' ')
print('Something went wrong with image '+jpgFileName+' SKIPPING #################################################')
#print(' ')
continue
if __name__ == '__main__':
with Pool(coresAmount) as p:
p.map(augmentAllImagesInPath, paths)