From 8529bd0c6fcfd6f13964d9219b4d797dfbb81690 Mon Sep 17 00:00:00 2001
From: Nazar Sturanec <nazarsturanec@gamil.com>
Date: Thu, 23 May 2024 17:27:36 +0300
Subject: [PATCH] Made some changes

---
 preprocessing.py |  3 +--
 train.py         | 69 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 2 deletions(-)
 create mode 100644 train.py

diff --git a/preprocessing.py b/preprocessing.py
index 25fa4ad..0a0fb40 100644
--- a/preprocessing.py
+++ b/preprocessing.py
@@ -42,7 +42,6 @@ def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
     transform = A.Compose([
         A.HorizontalFlip(p=0.5),
         A.RandomBrightnessContrast(p=0.2),
-        A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.5),
     ], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids']))
 
     bboxes = bboxes
@@ -56,7 +55,7 @@ def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
     return transformed_image, transformed_bboxes, transformed_category_ids
 
 def write_results(img_ann: ImageAnnotation):
-    for i in range(10):
+    for i in range(100):
         transformed_image, transformed_bboxes, transformed_category_ids, = image_processing(img_ann)
         cv2.imwrite(os.path.join(current_dataset_dir, images_dir, str(i)+ImageAnnotation(img_ann).image_path + '.jpg'), transformed_image)
         with open(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann).image_path + '.txt'), 'w') as f:
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..5086a0b
--- /dev/null
+++ b/train.py
@@ -0,0 +1,69 @@
+import math
+import os
+import datetime
+import cv2
+
+file_start_save = 'Zombobase-'+str(datetime.date.today())
+file_txt = []
+file_jpg = []
+labels_dir = 'labels'
+images_dir = 'images'
+current_dataset_dir = os.path.join('datasets', 'zombobase-current')
+
+
+def save(annotations, image, i, f_n):
+    cv2.imwrite(os.path.join(file_start_save, current_dataset_dir, f_n, file_jpg[i]), image)
+    with open(os.path.join(file_start_save, current_dataset_dir, f_n, file_txt[i]), 'w') as f:
+        for iii in range(len(annotations)):
+            f.write(annotations[iii])
+    annotations = []
+
+
+def piercing_photo_file():
+    for foldername, subfolders, filenames in os.walk(current_dataset_dir):
+        for subfolder in subfolders:
+            cu = os.path.join('datasets', 'zombobase-current',subfolder)
+            for foldername, subfolders, filenames in os.walk(cu):
+                for filename in filenames:
+                    f = filename.split('.')
+                    if f[-1] == 'txt':
+                        file_txt.append(filename)
+                    elif f[-1] == 'jpg':
+                        file_jpg.append(filename)
+
+
+def file_validation(annotations, file_txt, file_jpg):
+    os.makedirs(os.path.join(file_start_save, current_dataset_dir,'validation'))
+    for i in range(len(file_txt)):
+        image = cv2.imread(os.path.join(current_dataset_dir, images_dir, file_jpg[i]))
+        with open(os.path.join(current_dataset_dir, labels_dir, file_txt[i]), 'r') as file:
+            lines = file.readlines()
+            for line in lines:
+                annotations.append(line)
+        save(annotations, image, i, 'validation')
+
+def sort_file():
+    os.makedirs(os.path.join(file_start_save, current_dataset_dir))
+    annotations = []
+    folder_name = ['test', 'train']
+    percent_file = [0.20, 0.10]
+    for f_n, p_f in zip(folder_name, percent_file):
+        os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n))
+
+        for i in range(math.ceil(len(file_txt) * p_f)):
+            image = cv2.imread(os.path.join(current_dataset_dir, images_dir, file_jpg[i]))
+            with open(os.path.join(current_dataset_dir, labels_dir, file_txt[i]), 'r') as file:
+                lines = file.readlines()
+                for line in lines:
+                    annotations.append(line)
+                    save(annotations, image, i, f_n)
+                    file_txt.pop(i)
+                    file_jpg.pop(i)
+    file_validation(annotations, file_txt, file_jpg)
+
+
+def main():
+    piercing_photo_file()
+    sort_file()
+
+main()
\ No newline at end of file