From c172d3357d9630cdcd61c933e7e242de6bc6e811 Mon Sep 17 00:00:00 2001 From: hasslesstech Date: Wed, 3 Dec 2025 12:06:34 +0200 Subject: [PATCH] aug basic --- aug.py | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 aug.py diff --git a/aug.py b/aug.py new file mode 100644 index 0000000..3cc4571 --- /dev/null +++ b/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))