aug basic

aug.py

@@ -0,0 +1,64 @@
+import albumentations as a
+import numpy as np
+#from ultralytics import YOLO
+from os import listdir as ls
+import cv2 as cv
+#from pybboxes import BoundingBox as bb
+
+#import torch
+
+def b_c(b, s = (1920, 1080)):
+    return [
+        round(b[0] * s[0] - b[2] * s[0] / 2),
+        round(b[0] * s[0] + b[2] * s[0] / 2),
+        round(b[1] * s[1] - b[3] * s[1] / 2),
+        round(b[1] * s[1] + b[3] * s[1] / 2),
+    ]
+
+def sq(img, bbox):
+    for i in bbox:
+        b = b_c(i, s = (img.shape[1], img.shape[0]))
+        #print(b)
+        cv.rectangle(img, (b[0], b[2]), (b[1], b[3]), (255, 0, 0), 2)
+
+base_path = '../videos/total'
+tagt_path = '../videos/total-exp'
+
+t = a.Compose([
+    a.BBoxSafeRandomCrop(),
+    a.HorizontalFlip(p=0.5),
+    a.RandomBrightnessContrast(p=0.2)
+], bbox_params=a.BboxParams(format='yolo', label_fields=['class_labels']))
+
+fs = set([i.rsplit(".", 1)[0] for i in ls(base_path) if i != "classes.txt"])
+
+for i in list(fs)[:1]:
+    img = cv.cvtColor(cv.imread(f"{base_path}/{i}.jpg"), cv.COLOR_BGR2RGB)
+    box = [tuple(map(float, j.split()[1:])) for j in open(f"{base_path}/{i}.txt").read().split("\n") if j]
+    
+    #for j in open(f"{base_path}/{i}.txt").read().split("\n"):
+    #    box.append(tuple(map(float, j.split()[1:])))
+
+    print(box)
+
+    lbl = np.array(['ch'] * len(box))
+    box = np.array(box)
+
+    #sq(img, box)
+
+    r = t(image = img, bboxes = box, class_labels = lbl)
+
+    out_img = r['image']
+    out_box = r['bboxes']
+    out_lbl = r['class_labels']
+
+    sq(out_img, out_box)
+    
+    from matplotlib import pyplot as plt
+    plt.imshow(out_img)
+    plt.show()
+
+#y = YOLO("m.yaml")
+#y.info()
+
+#o = y.model(torch.randn(1, 3, 640, 640))