aug batchable

This commit is contained in:
ІО-23 Шмуляр Олег 2025-12-03 12:20:19 +02:00
parent c172d3357d
commit 830c3a0571

38
aug.py
View File

@ -1,11 +1,7 @@
import albumentations as a import albumentations as a
import numpy as np import numpy as np
#from ultralytics import YOLO
from os import listdir as ls from os import listdir as ls
import cv2 as cv import cv2 as cv
#from pybboxes import BoundingBox as bb
#import torch
def b_c(b, s = (1920, 1080)): def b_c(b, s = (1920, 1080)):
return [ return [
@ -24,6 +20,8 @@ def sq(img, bbox):
base_path = '../videos/total' base_path = '../videos/total'
tagt_path = '../videos/total-exp' tagt_path = '../videos/total-exp'
p = 1
t = a.Compose([ t = a.Compose([
a.BBoxSafeRandomCrop(), a.BBoxSafeRandomCrop(),
a.HorizontalFlip(p=0.5), a.HorizontalFlip(p=0.5),
@ -36,29 +34,23 @@ for i in list(fs)[:1]:
img = cv.cvtColor(cv.imread(f"{base_path}/{i}.jpg"), cv.COLOR_BGR2RGB) 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] 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)) lbl = np.array(['ch'] * len(box))
box = np.array(box) box = np.array(box)
#sq(img, box) for k in range(10):
r = t(image = img, bboxes = box, class_labels = lbl)
r = t(image = img, bboxes = box, class_labels = lbl) out_img = r['image']
out_box = r['bboxes']
out_lbl = r['class_labels']
out_img = r['image'] #print(f"open({tagt_path}/{p}-{i}-{k}.xml)")
out_box = r['bboxes'] #print("\n".join([" ".join(list(map(str, ['0'] + list(y)))) for y in out_box]) + "\n")
out_lbl = r['class_labels'] open(f"{tagt_path}/{p}-{i}-{k}.xml", 'w').write("\n".join([" ".join(list(map(str, ['0'] + list(y)))) for y in out_box]) + "\n")
cv.imwrite(f"{tagt_path}/{p}-{i}-{k}.jpg", cv.cvtColor(out_img, cv.COLOR_RGB2BGR))
sq(out_img, out_box) #sq(out_img, out_box)
from matplotlib import pyplot as plt #from matplotlib import pyplot as plt
plt.imshow(out_img) #plt.imshow(out_img)
plt.show() #plt.show()
#y = YOLO("m.yaml")
#y.info()
#o = y.model(torch.randn(1, 3, 640, 640))