Upload folder using huggingface_hub (part 4)

train/datasets/camera_obb/labels/val/4835158c093e3ffd52342a423809364754226b2b.txt

@@ -0,0 +1,2 @@
+0 0.2634105936686198 0.1935646791811343 0.2954395548502604 0.1943834997106481 0.2945770772298177 0.3010044578269676 0.2625481160481771 0.3001856372974537
+2 0.2549540456136068 0.06895374439380787 0.3694180043538411 0.06842640064380787 0.3703617095947266 0.7157927336516203 0.2558977508544922 0.7163200774016204

train/datasets/camera_obb/labels/val/53a01ff48d650b43cf17cc5fe1bf494aee25616c.txt

@@ -0,0 +1,3 @@
+0 0.3857922770182292 0.2871536327199348 0.5518802490234375 0.2918396719259108 0.5474978434244792 0.7824607957595897 0.3814098714192709 0.7777747565536137
+1 0.4630966389973958 0.3610527797897845 0.5125653889973958 0.3584376692206939 0.5143194580078125 0.463243638169709 0.4648507080078125 0.4658587487387996
+2 0.3695952555338541 0.9195484324089159 0.3863800048828125 0 0.9585496419270834 0.01745386372245318 0.9417649739583334 0.9520072756220379

train/datasets/camera_obb/labels/val/56c28acf11c5decc8d221f1ca2e3df6bd2d049d4.txt

@@ -0,0 +1,2 @@
+0 0.6469535827636719 0.007856623145532119 0.7689152956008911 0.01588265682907439 0.7596372365951538 0.3327922011189719 0.6376755237579346 0.3247661674354296
+1 0.6916757822036743 0.1198771464178887 0.6943172216415405 0.05089098258737589 0.7366292476654053 0.05453261658876567 0.7339878082275391 0.123518791589695

train/datasets/camera_obb/labels/val/5c0d510a70c11ec736e414189db16f45ddb5fc9d.txt

@@ -0,0 +1,4 @@
+0 0.7914837837219239 0.5136481391059028 0.8466664314270019 0.5115435282389323 0.8488089561462402 0.6890873379177518 0.7936263084411621 0.6911919487847222
+0 0.1983717083930969 0.2643807093302409 0.3806298494338989 0.2605639775594076 0.383868956565857 0.7494202507866754 0.2016108274459839 0.7532369825575087
+1 0.2537240982055664 0.3423541598849826 0.3012538909912109 0.3068993462456597 0.3272118091583252 0.4168797810872396 0.2796820163726806 0.4523345947265625
+2 0.7366658210754394 0.2537775039672852 0.9560970306396485 0.2454183154635959 0.9623207092285156 0.7617679172092013 0.7428895950317382 0.7701271057128907

train/datasets/camera_obb/labels/val/60768f20302a19d26b68b54686dd00207ec16ab1.txt

@@ -0,0 +1,2 @@
+1 0.6274074009486608 0.2644444329398019 0.6637037004743304 0.3488888876778739 0.6451851981026786 0.4922222137451172 0.6022222246442522 0.4033333369663784
+2 0.3651851835704985 0.9755555289132255 0.4014814830961682 0.04777777876172747 0.7362962995256697 0.192222227369036 0.6488888695126488 0.9922222682407924

train/datasets/camera_obb/labels/val/68044adadc1f99dd2d091ac5c8f107c595d3dfe4.txt

@@ -0,0 +1,3 @@
+0 0.6829468308833607 0.2528458712811403 0.7991093685752467 0.2554823805501048 0.7967449522854989 0.5849879795973089 0.6805824145936129 0.5823514703283443
+1 0.7417603074458607 0.3047691976038416 0.7694990726939419 0.3023852082012968 0.7712760523745888 0.3677826165781006 0.7435372871265077 0.3701666059806455
+2 0.4695597330729167 0.7637564767727428 0.4879421301055373 0.006196005667985508 0.8794279399671052 0.03624256687491621 0.8610455429344847 0.7938030082238445

train/datasets/camera_obb/labels/val/80fdea1b6d692db74f3904623dad229289fc654c.txt

@@ -0,0 +1,2 @@
+0 0.1839114129543304 0.5449801146529103 0.3442055881023407 0.5420456747972329 0.3473124802112579 0.9568405966722329 0.1870182752609253 0.9597750365279103
+2 0.1401884853839874 0.7399712715440124 0.1541937291622162 0 0.6873822212219238 0 0.7646484375 0.9709923664122138

train/datasets/camera_obb/labels/val/91dd2d6c11a95a8ab1fde7038ce8beeca2563f85.txt

@@ -0,0 +1,3 @@
+0 0.5491292953491211 0.005036898858532934 0.7265292485555013 0 0.7365162531534831 0.6719988246163923 0.5591162999471029 0.6857598036349176
+1 0.6448446273803711 0.07216061346545191 0.7101897557576498 0.08525005957323634 0.7039388656616211 0.2502416919091505 0.6385937372843424 0.237152245801366
+2 0.4242466608683268 0.9442284132906062 0.4336291631062826 0 0.7486366907755534 0 0.7392541885375976 0.9604403741345434

train/datasets/camera_obb/labels/val/94e8132cb5b3d9c7a7db2e4c01358c7ce0a1ba2d.txt

@@ -0,0 +1,3 @@
+0 0.43832794825236 0.4956274032592773 0.5842595100402832 0.4910593032836914 0.5907653172810873 0.8605449464586046 0.4448337554931641 0.8651130464341905
+1 0.4882832368214925 0.6242245038350424 0.4909009138743083 0.5528739823235406 0.5458494822184244 0.5564578904045953 0.5432318051656088 0.627808411916097
+2 0.3541666666666667 0.9565972222222222 0.1888020833333333 0.05034722222222222 0.81640625 0.5052083333333334 0.6666666666666666 0.9565972222222222

train/datasets/camera_obb/labels/val/a474b6a2458d7036c7b1e4858a320feb77cd1d15.txt

@@ -0,0 +1,3 @@
+0 0.4371829986572265 0.3666340611441721 0.6163790130615234 0.3680316860635554 0.6148435211181641 0.8115530962344629 0.4356475067138672 0.8101554713150797
+1 0.5193661117553711 0.438187704748329 0.5729095077514649 0.437737441644436 0.5732589340209961 0.5313378883347503 0.5197155380249023 0.5317881514386433
+2 0.1617824935913086 0.7529763513389717 0.178365592956543 0 0.7620236206054688 0 0.7454405212402344 0.7802020389635612

train/datasets/camera_obb/labels/val/a529352554ef426512417388797069bce2d7ae61.txt

@@ -0,0 +1,3 @@
+0 0.445787607828776 0.3697650372540509 0.5501012674967448 0.3702049424913195 0.5496566263834636 0.7034872775607639 0.4453429667154948 0.7030473723234953
+1 0.4792294311523437 0.4853821252893519 0.4809566243489583 0.4248517975983797 0.5169183349609375 0.4280949797453704 0.5151911417643229 0.4886253074363426
+2 0.2923527018229167 1 0.1658333333333333 0.3022222222222222 0.6783333333333333 0.07851851851851852 0.7097437540690105 1

train/datasets/camera_obb/labels/val/aca23d3803f305e35173fffd661249cbeb86f282.txt

@@ -0,0 +1,2 @@
+0 0.5129395167032877 0.07099106400101274 0.7655488967895507 0.06353409378616898 0.7728203455607097 0.8420400548864294 0.5202109654744466 0.8494970251012731
+1 0.6076282501220703 0.3633958604600694 0.6105879465738933 0.2010808874059606 0.7084142049153646 0.2067185013382523 0.7054545084635416 0.3690334743923611

train/datasets/camera_obb/labels/val/b1627e6fc85ec2548deaaaa92a440f0b33e2569b.txt

@@ -0,0 +1,2 @@
+0 0.4356682002544403 0.06801519791285197 0.7278732657432556 0.06890211502710979 0.7268820405006409 0.6495235761006674 0.4346769154071808 0.6486366589864095
+2 0.1624999940395355 0.5711111227671305 0.2783333361148834 0.1555555562178294 0.7056892514228821 0.07034533719221751 0.6966195702552795 0.5922128359476725

train/datasets/camera_obb/labels/val/c0a2be2294bc04ece2457da1ce96e0d8e0bbab34.txt

@@ -0,0 +1,3 @@
+0 0.4466235160827637 0.528122091293335 0.5309912681579589 0.5309156894683837 0.5293543179829915 0.9758451461791993 0.4449865659077962 0.973051643371582
+1 0.476533317565918 0.6769670963287353 0.4775744756062826 0.5851541042327881 0.5057493845621744 0.5880296230316162 0.5047082265218099 0.6798426151275635
+2 0.1989908377329508 0.9435964584350586 0.2131631374359131 0 0.5646679560343425 0.02329319417476654 0.5504956563313802 0.9899763107299805

train/datasets/camera_obb/labels/val/c1e1a508068bf406d5f8d5352888a7ddef501ea6.txt

@@ -0,0 +1,2 @@
+1 0.3092448276756084 0.2622222128965496 0.2800708061825913 0.04000000138856524 0.7312959654141317 0.06000000015428503 0.6496086728256361 0.2988888824432034
+0 0.09335693089307937 0.03666666568952817 0.9121750553097345 0.04000000138856524 0.760470013913855 0.9788888713346815 0.114751225024198 0.9644444603529448

train/datasets/camera_obb/labels/val/cb82349e33740e07c907a4de3efe22ea32558d52.txt

@@ -0,0 +1,3 @@
+0 0.4388638973236084 0.4919513702392578 0.58342604637146 0.4872781117757162 0.5902125835418701 0.8604966481526692 0.4456504344940185 0.865169906616211
+1 0.4875132560729981 0.6265329996744792 0.4890247344970703 0.5556292215983073 0.5451711654663086 0.5577570597330729 0.5436596870422363 0.6286608378092448
+2 0.2868754386901856 0.9090583165486653 0.1350000023841858 0.02666666706403097 0.7983333110809326 0.4422222137451172 0.7233333110809326 0.9622222264607747

train/datasets/camera_obb/labels/val/e3bfd3635dbc0efcab18a5f251a641039481d7c9.txt

@@ -0,0 +1,4 @@
+0 0.3613726615905762 0.4235392570495605 0.4915069580078125 0.4217455863952637 0.4937459945678711 0.7105350494384766 0.3636117458343506 0.7123287200927735
+0 0.7856846332550049 0.9481299718221029 0.825963020324707 0.9478384017944336 0.8262105941772461 1 0.785932207107544 1
+1 0.3884547710418701 0.4672733306884765 0.4202221393585205 0.4679524103800456 0.4195156574249267 0.526711908976237 0.3877482414245605 0.526032829284668
+2 0.283979082107544 0.795031992594401 0.3017405986785889 0.1728111267089844 0.8135955810546875 0.1987864176432292 0.7958340167999267 0.8210072835286458

train/datasets/camera_obb/labels/val/e65982052a28bb225c41129bf520b636d3f5c993.txt

@@ -0,0 +1,2 @@
+0 0.3148372904459635 0.09865843757750496 0.560060780843099 0.09010159931485615 0.5677359517415365 0.8881268213665675 0.322512461344401 0.8966835627480159
+1 0.3833293151855469 0.3840068756587922 0.3866056823730469 0.2127370440770709 0.4801041158040364 0.2192262437608507 0.4768277486165365 0.3904960753425719

train/datasets/camera_obb/labels/val/ecd391135d962276d9695a3b92ca5bc99ab47472.txt

@@ -0,0 +1 @@
+0 0.5728878021240235 0.430283440483941 0.6786713600158691 0.4347163730197482 0.6739102363586426 0.7937912835015191 0.5681266784667969 0.7893583509657118

train/datasets/camera_obb/labels/val/eee25dbb6e5acfd232f3b06d2634622146508cff.txt

@@ -0,0 +1,3 @@
+0 0.5491292953491211 0.005036898858532934 0.7265292485555013 0 0.7365162531534831 0.6719988246163923 0.5591162999471029 0.6857598036349176
+1 0.6448446273803711 0.07216061346545191 0.7101897557576498 0.08525005957323634 0.7039388656616211 0.2502416919091505 0.6385937372843424 0.237152245801366
+2 0.4242466608683268 0.9442284132906062 0.4336291631062826 0 0.7486366907755534 0 0.7392541885375976 0.9604403741345434

train/datasets/camera_obb/labels/val/f292d8e2973ceb973a7b69a22289e8615128ec1a.txt

@@ -0,0 +1,3 @@
+0 0.5153091152106047 0.5150942690737612 0.5948823692832467 0.5121440165751689 0.5971620880807873 0.7062115093609235 0.5175888340081453 0.7091617618595157
+1 0.5268022951199607 0.6156268695453265 0.5286383019499556 0.5505234726914414 0.5675753889412716 0.5539892798071508 0.5657393821112766 0.6190926766610361
+2 0.3210640804762773 0.8932733209283502 0.3354901667782549 0 0.8460694597410624 0 0.8316433734390847 0.918549023877393

train/datasets/camera_obb/labels/val/f927b4372f85162d0bad7059228d3490ffb48f79.txt

@@ -0,0 +1,3 @@
+0 0.4444444675901237 0.2533333396911621 0.5142207227062446 0.2511111068725586 0.5036025844743227 0.33 0.4535456842936754 0.3266666793823242
+1 0.4629607737674648 0.2924523162841797 0.464551645741121 0.2645793914794922 0.5033547984048368 0.2657676696777344 0.5017639264311807 0.2936405944824219
+2 0.2411971889665102 0.4442596054077149 0.2634520026197206 0.03292970418930054 0.7079076148544156 0.04583231449127197 0.6856528012012052 0.4571622085571289

train/datasets/camera_obb/labels/val/fc0f3cdc42d1af9b05a1ff82f53147d3b9133b1c.txt

@@ -0,0 +1,3 @@
+0 0.3746675252914429 0.5141487121582031 0.4425457715988159 0.5148464441299438 0.4416691859563192 0.7067247629165649 0.3737909396489461 0.7060270309448242
+1 0.384706695874532 0.5441105365753174 0.407295028368632 0.543700098991394 0.4076470136642456 0.5872790813446045 0.3850586811701457 0.5876895189285278
+2 0.3042814135551453 1 0.3256328304608663 0.2897551953792572 0.6319967110951742 0.3097002804279327 0.6106452941894531 1

train/datasets/camera_obb/labels/val/ff05ed7bd591f1f3c07d87d9a50491eb5c28d36c.txt

@@ -0,0 +1 @@
+0 0.4927224245938388 0.5001363153148957 0.5120611364191229 0.500320252608603 0.5118539636785334 0.5492714877022998 0.4925152518532493 0.5490875504085925

train/requirements.txt

@@ -0,0 +1,6 @@
+# Training
+ultralytics
+torch
+ddgs
+numpy
+opencv-python

train/testObbOpenCv.py

@@ -0,0 +1,311 @@
+import argparse
+import math
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+REPO_DIR = Path(__file__).parent.parent
+DEFAULT_MODEL = REPO_DIR / "model" / "best.onnx"
+
+CLASSES = {
+    0: "Body",
+    1: "Lens",
+    2: "System",
+}
+
+COLORS = {
+    0: (76, 175, 80),
+    1: (33, 150, 243),
+    2: (255, 152, 0),
+}
+
+
+def letterbox(
+    image: np.ndarray,
+    imageSize: int,
+    fillColor: tuple[int, int, int] = (114, 114, 114),
+) -> tuple[np.ndarray, float, int, int]:
+    height, width = image.shape[:2]
+    scale = min(imageSize / width, imageSize / height)
+
+    resizedWidth = int(round(width * scale))
+    resizedHeight = int(round(height * scale))
+
+    resized = cv2.resize(
+        image,
+        (resizedWidth, resizedHeight),
+        interpolation=cv2.INTER_LINEAR,
+    )
+
+    padX = (imageSize - resizedWidth) // 2
+    padY = (imageSize - resizedHeight) // 2
+
+    canvas = np.full(
+        (imageSize, imageSize, 3),
+        fillColor,
+        dtype=np.uint8,
+    )
+    canvas[
+        padY : padY + resizedHeight,
+        padX : padX + resizedWidth,
+    ] = resized
+
+    return canvas, scale, padX, padY
+
+
+def parseOutput(
+    output: np.ndarray,
+    confidenceThreshold: float,
+    scale: float,
+    padX: int,
+    padY: int,
+    imageShape: tuple[int, int, int],
+) -> tuple[
+    list[tuple[tuple[float, float], tuple[float, float], float]], list[float], list[int]
+]:
+    if output.ndim == 3:
+        output = output[0]
+
+    if output.shape[0] < output.shape[1]:
+        predictions = output.T
+    else:
+        predictions = output
+
+    classCount = len(CLASSES)
+    expectedValues = 5 + classCount
+    if predictions.shape[1] != expectedValues:
+        raise ValueError(
+            f"Expected {expectedValues} ONNX outputs per prediction, "
+            f"got {predictions.shape[1]}"
+        )
+
+    imageHeight, imageWidth = imageShape[:2]
+    boxes = []
+    scores = []
+    classIds = []
+
+    for prediction in predictions:
+        classScores = prediction[4 : 4 + classCount]
+        classId = int(np.argmax(classScores))
+        score = float(classScores[classId])
+        if score < confidenceThreshold:
+            continue
+
+        centerX, centerY, width, height = prediction[:4]
+        angleRadians = float(prediction[4 + classCount])
+
+        centerX = (float(centerX) - padX) / scale
+        centerY = (float(centerY) - padY) / scale
+        width = float(width) / scale
+        height = float(height) / scale
+
+        centerX = min(max(centerX, 0.0), float(imageWidth - 1))
+        centerY = min(max(centerY, 0.0), float(imageHeight - 1))
+        width = max(width, 1.0)
+        height = max(height, 1.0)
+
+        boxes.append(
+            (
+                (centerX, centerY),
+                (width, height),
+                math.degrees(angleRadians),
+            )
+        )
+        scores.append(score)
+        classIds.append(classId)
+
+    return boxes, scores, classIds
+
+
+def runRotatedNms(
+    boxes: list[tuple[tuple[float, float], tuple[float, float], float]],
+    scores: list[float],
+    classIds: list[int],
+    confidenceThreshold: float,
+    nmsThreshold: float,
+) -> list[int]:
+    keptIndexes = []
+
+    for classId in sorted(set(classIds)):
+        localIndexes = [
+            index for index, boxClassId in enumerate(classIds) if boxClassId == classId
+        ]
+        localBoxes = [boxes[index] for index in localIndexes]
+        localScores = [scores[index] for index in localIndexes]
+
+        selected = cv2.dnn.NMSBoxesRotated(
+            localBoxes,
+            localScores,
+            confidenceThreshold,
+            nmsThreshold,
+        )
+
+        if len(selected) == 0:
+            continue
+
+        for selectedIndex in np.array(selected).flatten():
+            keptIndexes.append(localIndexes[int(selectedIndex)])
+
+    return keptIndexes
+
+
+def drawDetections(
+    image: np.ndarray,
+    boxes: list[tuple[tuple[float, float], tuple[float, float], float]],
+    scores: list[float],
+    classIds: list[int],
+    indexes: list[int],
+) -> np.ndarray:
+    output = image.copy()
+
+    for index in indexes:
+        classId = classIds[index]
+        color = COLORS.get(classId, (255, 255, 255))
+        label = f"{CLASSES.get(classId, classId)} {scores[index]:.2f}"
+
+        points = cv2.boxPoints(boxes[index])
+        points = np.intp(points)
+
+        cv2.polylines(output, [points], True, color, 2, cv2.LINE_AA)
+
+        labelX = int(points[:, 0].min())
+        labelY = int(points[:, 1].min()) - 8
+        labelY = max(labelY, 20)
+
+        textSize, baseline = cv2.getTextSize(
+            label,
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.6,
+            2,
+        )
+        textWidth, textHeight = textSize
+
+        cv2.rectangle(
+            output,
+            (labelX, labelY - textHeight - baseline),
+            (labelX + textWidth + 6, labelY + baseline),
+            color,
+            -1,
+        )
+        cv2.putText(
+            output,
+            label,
+            (labelX + 3, labelY),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.6,
+            (255, 255, 255),
+            2,
+            cv2.LINE_AA,
+        )
+
+    return output
+
+
+def getOutputPath(imagePath: Path, outputPath: Path | None) -> Path:
+    if outputPath:
+        return outputPath
+
+    return imagePath.with_name(f"{imagePath.stem}_obb{imagePath.suffix}")
+
+
+def parseArgs() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("image", type=Path, help="Image to run inference on.")
+    parser.add_argument(
+        "--model",
+        type=Path,
+        default=DEFAULT_MODEL,
+        help="Path to the exported YOLO OBB ONNX model.",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=None,
+        help="Path to save the image with bounding box overlays.",
+    )
+    parser.add_argument(
+        "--imgsz",
+        type=int,
+        default=960,
+        help="Input image size used when exporting the ONNX model.",
+    )
+    parser.add_argument(
+        "--conf",
+        type=float,
+        default=0.25,
+        help="Confidence threshold for detections.",
+    )
+    parser.add_argument(
+        "--iou",
+        type=float,
+        default=0.45,
+        help="Rotated NMS IoU threshold.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parseArgs()
+    imagePath = args.image
+    modelPath = args.model
+    outputPath = getOutputPath(imagePath, args.output)
+
+    if not imagePath.exists():
+        raise FileNotFoundError(f"Image not found: {imagePath}")
+
+    if not modelPath.exists():
+        raise FileNotFoundError(f"ONNX model not found: {modelPath}")
+
+    image = cv2.imread(str(imagePath))
+    if image is None:
+        raise ValueError(f"Could not read image: {imagePath}")
+
+    inputImage, scale, padX, padY = letterbox(image, args.imgsz)
+    blob = cv2.dnn.blobFromImage(
+        inputImage,
+        scalefactor=1 / 255.0,
+        size=(args.imgsz, args.imgsz),
+        mean=(0, 0, 0),
+        swapRB=True,
+        crop=False,
+    )
+
+    net = cv2.dnn.readNetFromONNX(str(modelPath))
+    net.setInput(blob)
+    output = net.forward()
+
+    boxes, scores, classIds = parseOutput(
+        output,
+        args.conf,
+        scale,
+        padX,
+        padY,
+        image.shape,
+    )
+    keptIndexes = runRotatedNms(
+        boxes,
+        scores,
+        classIds,
+        args.conf,
+        args.iou,
+    )
+
+    result = drawDetections(
+        image,
+        boxes,
+        scores,
+        classIds,
+        keptIndexes,
+    )
+
+    outputPath.parent.mkdir(parents=True, exist_ok=True)
+    if not cv2.imwrite(str(outputPath), result):
+        raise ValueError(f"Could not write output image: {outputPath}")
+
+    print(f"Detections: {len(keptIndexes)}")
+    print(f"Saved: {outputPath}")
+
+
+if __name__ == "__main__":
+    main()

train/trainObbModel.py

@@ -0,0 +1,197 @@
+import argparse
+import random
+import shutil
+from pathlib import Path
+
+import torch
+import yaml
+from ultralytics import YOLO
+
+classes = {
+    0: "Body",
+    1: "Lens",
+    2: "System",
+}
+
+# ~/Outflock
+REPO_DIR = Path(__file__).parent.parent
+IMAGES_DIR = REPO_DIR / "train/data/images"
+LABELS_DIR = REPO_DIR / "train/data/labels"
+DATASET_DIR = REPO_DIR / "train/datasets/camera_obb"
+RUNS_DIR = REPO_DIR / "train/runs"
+
+ONNX_DIR = REPO_DIR / "model"
+
+valRatio = 0.2
+seed = 42
+
+
+def hasValidLabel(imagePath: Path) -> bool:
+    labelPath = LABELS_DIR / f"{imagePath.stem}.txt"
+    if not labelPath.exists():
+        return False
+
+    lines = labelPath.read_text().strip().splitlines()
+    if not lines:
+        return False
+
+    for line in lines:
+        parts = line.split()
+        if len(parts) != 9:
+            return False
+
+        classId = int(parts[0])
+        if classId not in classes:
+            return False
+
+        coords = [float(value) for value in parts[1:]]
+        if any(value < 0 or value > 1 for value in coords):
+            return False
+
+    return True
+
+
+def copyExample(imagePath: Path, split: str) -> None:
+    labelPath = LABELS_DIR / f"{imagePath.stem}.txt"
+
+    imageOut = DATASET_DIR / "images" / split / imagePath.name
+    labelOut = DATASET_DIR / "labels" / split / labelPath.name
+
+    imageOut.parent.mkdir(parents=True, exist_ok=True)
+    labelOut.parent.mkdir(parents=True, exist_ok=True)
+
+    shutil.copy2(imagePath, imageOut)
+    shutil.copy2(labelPath, labelOut)
+
+
+def prepareDataset() -> Path:
+    if DATASET_DIR.exists():
+        shutil.rmtree(DATASET_DIR)
+
+    imagePaths = sorted(
+        path
+        for path in IMAGES_DIR.iterdir()
+        if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}
+        and hasValidLabel(path)
+    )
+
+    random.Random(seed).shuffle(imagePaths)
+
+    valCount = max(1, int(len(imagePaths) * valRatio))
+    valImages = set(imagePaths[:valCount])
+    trainImages = imagePaths[valCount:]
+
+    for imagePath in trainImages:
+        copyExample(imagePath, "train")
+
+    for imagePath in valImages:
+        copyExample(imagePath, "val")
+
+    dataYaml = DATASET_DIR / "data.yaml"
+    dataYaml.write_text(
+        yaml.safe_dump(
+            {
+                "path": str(DATASET_DIR.resolve()),
+                "train": "images/train",
+                "val": "images/val",
+                "names": classes,
+            },
+            sort_keys=False,
+        )
+    )
+
+    print(f"Prepared {len(trainImages)} train and {len(valImages)} val images")
+    return dataYaml
+
+
+def trainModel(dataYaml: Path, onnx: bool = False) -> None:
+    print(f"CUDA available: {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+
+    model = YOLO("yolov8m-obb.pt")
+    model.train(
+        data=str(dataYaml),
+        epochs=50,
+        imgsz=960,
+        project=str(RUNS_DIR),
+        name="flockOBB",
+        task="obb",
+        batch=16,
+        device=0,
+        workers=8,
+        patience=20,
+        pretrained=True,
+        optimizer="auto",
+        amp=True,
+        # Detection-specific defaults worth making explicit.
+        single_cls=False,
+        rect=False,
+        cache=False,
+        # Augmentation. Conservative for real camera detection.
+        degrees=5,
+        translate=0.08,
+        scale=0.4,
+        shear=0.0,
+        perspective=0.0005,
+        flipud=0.0,
+        fliplr=0.5,
+        mosaic=0.7,
+        mixup=0.05,
+        copy_paste=0.0,
+    )
+
+    if not onnx:
+        return
+
+    # Optional: export to ONNX and weights for OpenCV later.
+    bestWeights = Path(model.trainer.best)
+    if not bestWeights.exists():
+        saveDir = Path(model.trainer.save_dir)
+        bestWeights = saveDir / "weights" / "best.pt"
+
+    if not bestWeights.exists():
+        raise FileNotFoundError(f"Could not find trained best weights at {bestWeights}")
+
+    print(f"Best weights saved to: {bestWeights}")
+
+    exportModel = YOLO(str(bestWeights))
+    onnxPath = Path(
+        exportModel.export(
+            format="onnx",
+            imgsz=960,
+            opset=12,
+            simplify=True,
+            dynamic=False,
+        )
+    )
+
+    ONNX_DIR.mkdir(parents=True, exist_ok=True)
+    targetPath = ONNX_DIR / onnxPath.name
+    shutil.copy2(onnxPath, targetPath)
+    print(f"ONNX model copied to: {targetPath}")
+
+
+def parseArgs() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--onnx",
+        action="store_true",
+        help="Export the trained best weights to ONNX after training.",
+    )
+    parser.add_argument(
+        "--clean",
+        action="store_true",
+        help="Clean Non-ONNX model directories.",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parseArgs()
+    dataYaml = prepareDataset()
+    trainModel(dataYaml, onnx=args.onnx)
+
+    if args.clean:
+        shutil.rmtree(RUNS_DIR)
+        print(f"Removed Non-ONNX model directory: {RUNS_DIR}")