65 lines
1.6 KiB
Python
65 lines
1.6 KiB
Python
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))
|