From 41e1571c4d855254b43e46d9cc0de112a4adeba9 Mon Sep 17 00:00:00 2001
From: hasslesstech <hasslesstech@tutanota.com>
Date: Thu, 4 Dec 2025 13:19:41 +0200
Subject: [PATCH] update upstaged changes

---
 aug.py            | 12 ++++++++----
 detect.py         | 10 ++++++++++
 live.py           | 24 ++++++++++++++++++++++++
 show-box.py       | 33 +++++++++++++++++++++++++++++++++
 train.py          |  4 ++++
 video-cruncher.py | 31 +++++++++++++++++++++++++++++++
 6 files changed, 110 insertions(+), 4 deletions(-)
 create mode 100644 detect.py
 create mode 100644 live.py
 create mode 100644 show-box.py
 create mode 100644 train.py
 create mode 100644 video-cruncher.py

diff --git a/aug.py b/aug.py
index 444c0fd..5ae74b8 100644
--- a/aug.py
+++ b/aug.py
@@ -23,14 +23,18 @@ tagt_path = '../videos/total-exp'
 p = 1
 
 t = a.Compose([
+    a.SafeRotate(limit=80, p=0.8),
     a.BBoxSafeRandomCrop(),
     a.HorizontalFlip(p=0.5),
-    a.RandomBrightnessContrast(p=0.2)
+    a.VerticalFlip(p=0.5),
+    a.RandomBrightnessContrast(p=0.9, brightness_limit = (-0.2, 0.7), contrast_limit = (-0.6, 0.4)),
+    a.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.8),
+    a.ChromaticAberration(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]:
+for i in list(fs):
     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]
     
@@ -44,9 +48,9 @@ for i in list(fs)[:1]:
         out_box = r['bboxes']
         out_lbl = r['class_labels']
 
-        #print(f"open({tagt_path}/{p}-{i}-{k}.xml)")
+        print(f"open({tagt_path}/{p}-{i}-{k}.txt)")
         #print("\n".join([" ".join(list(map(str, ['0'] + list(y)))) for y in out_box]) + "\n")
-        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")
+        open(f"{tagt_path}/{p}-{i}-{k}.txt", '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)
diff --git a/detect.py b/detect.py
new file mode 100644
index 0000000..2ec0fca
--- /dev/null
+++ b/detect.py
@@ -0,0 +1,10 @@
+from ultralytics import YOLO
+from matplotlib import pyplot as plt
+
+m = YOLO("best.pt")
+
+res = m("../dataset/images/train/1-1-0000001-1.jpg")
+res[0].show()
+#print(res)
+#plt.imshow(res)
+#plt.show()
diff --git a/live.py b/live.py
new file mode 100644
index 0000000..695908e
--- /dev/null
+++ b/live.py
@@ -0,0 +1,24 @@
+from ultralytics import YOLO
+import json
+import cv2 as cv
+
+m = YOLO("best.pt")
+
+c = cv.VideoCapture(0)
+
+while True:
+    _, frame = c.read()
+
+    res = m(frame)
+    for i in json.loads(res[0].to_json()):
+        i = i['box']
+        cv.rectangle(frame, (round(i['x1']), round(i['y1'])), (round(i['x2']), round(i['y2'])), (0, 0, 255), 3)
+        #print(i)
+
+    cv.imshow('frame', frame)
+
+    if cv.waitKey(1) == ord('q'):
+        break
+
+c.release()
+cv.destroyAllWindows()
diff --git a/show-box.py b/show-box.py
new file mode 100644
index 0000000..0da7306
--- /dev/null
+++ b/show-box.py
@@ -0,0 +1,33 @@
+import albumentations as a
+import numpy as np
+from os import listdir as ls
+import cv2 as cv
+from sys import argv, exit
+
+if len(argv) != 2:
+    exit(1)
+
+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)
+
+img = cv.cvtColor(cv.imread(f"{argv[1]}.jpg"), cv.COLOR_BGR2RGB)
+box = [tuple(map(float, j.split()[1:])) for j in open(f"{argv[1]}.txt").read().split("\n") if j]
+
+box = np.array(box)
+
+sq(img, box)
+
+from matplotlib import pyplot as plt
+plt.imshow(img)
+plt.show()
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..123c4bb
--- /dev/null
+++ b/train.py
@@ -0,0 +1,4 @@
+from ultralytics import YOLO
+
+m = YOLO("m.yaml")
+m.train(data = "../dataset/d.yaml", epochs = 100, imgsz = 300)
diff --git a/video-cruncher.py b/video-cruncher.py
new file mode 100644
index 0000000..e876d8b
--- /dev/null
+++ b/video-cruncher.py
@@ -0,0 +1,31 @@
+from ultralytics import YOLO
+import json
+import cv2 as cv
+from sys import argv, exit
+
+if len(argv) != 2:
+    exit(1)
+
+m = YOLO("best.pt")
+
+ci = cv.VideoCapture(argv[1])
+target_dir = argv[1].rsplit(".", 1)[0] + "-p"
+
+k = 0
+
+while True:
+    ret, frame = ci.read()
+
+    if not ret:
+        break
+
+    res = m(frame)
+
+    for i in json.loads(res[0].to_json()):
+        i = i['box']
+        cv.rectangle(frame, (round(i['x1']), round(i['y1'])), (round(i['x2']), round(i['y2'])), (0, 0, 255), 3)
+
+    cv.imwrite(f"{target_dir}/{k:07d}.jpg", frame)
+    k += 1
+
+ci.release()