From ed1af3de970cb950a780dbe2e464b86603657481 Mon Sep 17 00:00:00 2001
From: Nazar Sturanec <nazarsturanec@gamil.com>
Date: Wed, 22 May 2024 15:55:31 +0300
Subject: [PATCH] Made some changes

---
 preprocessing.py | 90 +++++++++++++++++++++++++++++-------------------
 1 file changed, 54 insertions(+), 36 deletions(-)

diff --git a/preprocessing.py b/preprocessing.py
index 13538d7..25fa4ad 100644
--- a/preprocessing.py
+++ b/preprocessing.py
@@ -1,20 +1,15 @@
 import os.path
-import time
-
+import albumentations as A
 import cv2
-import albumentations as alb
-from os import listdir
-from os.path import isfile, join
 from pathlib import Path
 
 labels_dir = 'labels'
 images_dir = 'images'
 current_dataset_dir = os.path.join('datasets', 'zombobase-current')
-
 class ImageAnnotation:
 
     def read_annotations(self) -> [[]]:
-        with open(self.annotation_path, 'r') as f:
+        with open(self.dataset_annotation_path, 'r') as f:
             rows = f.readlines()
             arr = []
             for row in rows:
@@ -23,50 +18,73 @@ class ImageAnnotation:
                 coordinates = [float(n) for n in str_coordinates]
                 coordinates.append(class_num)
                 arr.append(coordinates)
+
             return arr
 
     def __init__(self, image_path):
         self.image_path = image_path
         self.image_name = Path(image_path).stem
-        self.dataset_image_path = os.path.join(current_dataset_dir, images_dir, self.image_name, os.path.basename(image_path))
-        self.image = cv2.imread(image_path)
+        self.dataset_image_path = os.path.join(current_dataset_dir, images_dir, self.image_path + '.jpg')
+        self.image = cv2.imread(self.dataset_image_path)
 
-        self.annotation_path = os.path.join(labels_dir, self.image_name, '.txt')
-        self.dataset_annotation_path = os.path.join(current_dataset_dir, labels_dir, self.image_name, '.txt')
+        self.annotation_path = os.path.join(labels_dir, self.image_path + '.txt')
+        self.dataset_annotation_path = os.path.join(current_dataset_dir, labels_dir, self.image_path + '.txt')
         self.annotations = self.read_annotations()
 
+
 def image_processing(img_ann: ImageAnnotation) -> [ImageAnnotation]:
-    # return structure example:
-    # utilize transform albumentations here
-    return [ImageAnnotation(f'{img_ann.image_name}1', image1, bboxes1 ),
-            ImageAnnotation(f'{img_ann.image_name}2', image2, bboxes2),
-            ...
-                ]
+    category_ids = []
+    bboxes = ImageAnnotation(img_ann).read_annotations()
+    for i in range(len(bboxes)):
+        category_ids.append(bboxes[i][4])
+        bboxes[i].pop(4)
+
+    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
+    imag = ImageAnnotation(img_ann).image
+
+    transformed = transform(image=imag, bboxes=bboxes, category_ids=category_ids)
+    transformed_image = transformed['image']
+    transformed_bboxes = transformed['bboxes']
+    transformed_category_ids = transformed['category_ids']
+
+    return transformed_image, transformed_bboxes, transformed_category_ids
 
 def write_results(img_ann: ImageAnnotation):
-    # write image cv2.imwrite(, image) dataset_image_path
-    # write img_ann.annotations into new file with name dataset_annotation_path
+    for i in range(10):
+        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:
+            for bbox, category_id in zip(transformed_bboxes, transformed_category_ids):
+                x_center, y_center, width, height = bbox
+                cla = category_id
+                f.write(f"{cla} {x_center} {y_center} {width} {height}\n")
+#
+#
+def process_image():
+    file_annotation = []
+    file_annotation_finished =[]
+    for foldername, subfolders, filenames in os.walk(os.path.join(current_dataset_dir,images_dir)):
+        file_annotation.append(filenames)
+        print(file_annotation)
 
+    for i in range(len(file_annotation[0])):
+        ff = file_annotation[0][i].split('.')
+        ff.pop(-1)
+        ff = '.'.join(ff)
+        file_annotation_finished.append(ff)
 
-def process_image(img_ann):
-    results = image_processing(img_ann)
-    for res_ann in results:
-        write_results(res_ann)
-    write_results(img_ann)
-    os.remove(img_ann.image_path)
-    os.remove(img_ann.annotation_path)
-
+    return file_annotation_finished
 
 def main():
-    while True:
-        images = os.listdir(images_dir)
-        if len(images) == 0:
-            time.sleep(5)
-            continue
-
-        for image in images:
-            image_path = os.path.join(images_dir, image)
-            process_image(ImageAnnotation(image_path))
+    process_image()
+    for i in process_image():
+        write_results(i)
 
 if __name__ == '__main__':
     main()
\ No newline at end of file