456 lines
27 KiB
Plaintext
456 lines
27 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"id": "cd1c6f50-4456-4853-b76c-1346fe9476f8",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"-> Detected GPU: NVIDIA GeForce GTX 1050 Ti\n",
|
|
"-> GTX 1050 detected! Force-disabling cuDNN to prevent crashes...\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"import torch\n",
|
|
"\n",
|
|
"# === STEP 1: DYNAMIC HARDWARE CHECK ===\n",
|
|
"is_1050 = False\n",
|
|
"\n",
|
|
"if torch.cuda.is_available():\n",
|
|
" # Retrieve the exact GPU name (e.g., \"NVIDIA GeForce GTX 1050 Ti\")\n",
|
|
" gpu_name = torch.cuda.get_device_name(0)\n",
|
|
" print(f\"-> Detected GPU: {gpu_name}\")\n",
|
|
" \n",
|
|
" # Check if this is a 1050 series card\n",
|
|
" if \"1050\" in gpu_name:\n",
|
|
" is_1050 = True\n",
|
|
" print(\"-> GTX 1050 detected! Force-disabling cuDNN to prevent crashes...\")\n",
|
|
" \n",
|
|
" # Disable cuDNN entirely for PyTorch\n",
|
|
" torch.backends.cudnn.enabled = False\n",
|
|
" torch.backends.cudnn.benchmark = False\n",
|
|
" else:\n",
|
|
" print(\"-> Modern GPU detected. Keeping cuDNN enabled for maximum speed.\")\n",
|
|
"else:\n",
|
|
" print(\"-> No CUDA GPU detected. Running on CPU.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"id": "ee6d3405-39e3-4109-84e0-f1ba6c7df74b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%matplotlib inline\n",
|
|
"\n",
|
|
"import os\n",
|
|
"import random\n",
|
|
"import cv2\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"from ultralytics import YOLO\n",
|
|
"\n",
|
|
"# --- CONFIGURATION ---\n",
|
|
"DATASET_DIR = os.path.abspath(\"dataset_v5\") # Absolute path to your dataset folder\n",
|
|
"# VIDEO_PATH = \"Screen_Recording_20260710_101822_Hades' Star.mp4\" # Path to your test video\n",
|
|
"VIDEO_PATH = \"Screen_Recording_20260710_142109_Hades' Star.mp4\"\n",
|
|
"\n",
|
|
"# The exact list from your label.py script\n",
|
|
"CLASS_NAMES = [\n",
|
|
" \"guardian\",\n",
|
|
" \"interceptor\",\n",
|
|
" \"sentinel\",\n",
|
|
" \"rocket\",\n",
|
|
" \"bomber\",\n",
|
|
"]\n",
|
|
"# ---------------------"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"id": "0245800c-f39d-479d-8093-a8573223d975",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"-> Created dataset.yaml successfully.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# === STEP 1: CREATE THE YOLO CONFIGURATION FILE ===\n",
|
|
"yaml_content = f\"\"\"\n",
|
|
"path: {DATASET_DIR}\n",
|
|
"train: images # path relative to 'path'\n",
|
|
"val: images # for a quick test, we will validate on the same images\n",
|
|
"\n",
|
|
"names:\n",
|
|
"\"\"\"\n",
|
|
"for idx, name in enumerate(CLASS_NAMES):\n",
|
|
" yaml_content += f\" {idx}: {name}\\n\"\n",
|
|
"\n",
|
|
"yaml_path = \"dataset.yaml\"\n",
|
|
"with open(yaml_path, \"w\") as f:\n",
|
|
" f.write(yaml_content.strip())\n",
|
|
"\n",
|
|
"print(f\"-> Created {yaml_path} successfully.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a3e69197-7565-4464-8381-ffc871818819",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\n",
|
|
"-> Initializing YOLO training...\n",
|
|
"New https://pypi.org/project/ultralytics/8.4.96 available 😃 Update with 'pip install -U ultralytics'\n",
|
|
"Ultralytics 8.4.92 🚀 Python-3.10.12 torch-2.5.1+cu124 CUDA:0 (NVIDIA GeForce GTX 1050 Ti, 4030MiB)\n",
|
|
"\u001b[34m\u001b[1mengine/trainer: \u001b[0magnostic_nms=False, amp=False, angle=1.0, augment=False, auto_augment=randaugment, batch=4, bgr=0.0, box=7.5, cache=False, cfg=None, classes=None, close_mosaic=10, cls=0.5, cls_pw=0.0, cls_remap=True, compile=False, conf=None, copy_paste=0.0, copy_paste_mode=flip, cos_lr=False, cutmix=0.0, data=dataset.yaml, degrees=0.0, deterministic=True, device=, dfl=1.5, dis=6.0, distill_model=None, dnn=False, dropout=0.0, dynamic=False, embed=None, end2end=None, epochs=50, erasing=0.4, exist_ok=False, fliplr=0.5, flipud=0.0, format=torchscript, fraction=1.0, freeze=None, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, imgsz=640, iou=0.7, keras=False, kobj=1.0, line_width=None, lr0=0.01, lrf=0.01, mask_ratio=4, max_det=300, mixup=0.0, mode=train, model=yolov8n.pt, momentum=0.937, mosaic=1.0, multi_scale=0.0, name=train-22, nbs=64, nms=False, opset=None, optimize=False, optimizer=auto, overlap_mask=True, patience=5, perspective=0.0, plots=True, pose=12.0, pretrained=True, profile=False, project=None, quantize=None, rect=False, resume=False, retina_masks=False, rle=1.0, save=True, save_conf=False, save_crop=False, save_dir=/home/nathan/HadesStarYOLO/runs/detect/train-22, save_frames=False, save_json=False, save_period=-1, save_txt=False, scale=0.5, seed=0, shear=0.0, show=False, show_boxes=True, show_conf=True, show_labels=True, simplify=True, single_cls=False, source=None, split=val, stream_buffer=False, task=detect, time=None, tracker=tracktrack.yaml, translate=0.1, val=True, verbose=True, vid_stride=1, visualize=False, warmup_bias_lr=0.1, warmup_epochs=3.0, warmup_momentum=0.8, weight_decay=0.0005, workers=8, workspace=None\n",
|
|
"Overriding model.yaml nc=80 with nc=5\n",
|
|
"\n",
|
|
" from n params module arguments \n",
|
|
" 0 -1 1 464 ultralytics.nn.modules.conv.Conv [3, 16, 3, 2] \n",
|
|
" 1 -1 1 4672 ultralytics.nn.modules.conv.Conv [16, 32, 3, 2] \n",
|
|
" 2 -1 1 7360 ultralytics.nn.modules.block.C2f [32, 32, 1, True] \n",
|
|
" 3 -1 1 18560 ultralytics.nn.modules.conv.Conv [32, 64, 3, 2] \n",
|
|
" 4 -1 2 49664 ultralytics.nn.modules.block.C2f [64, 64, 2, True] \n",
|
|
" 5 -1 1 73984 ultralytics.nn.modules.conv.Conv [64, 128, 3, 2] \n",
|
|
" 6 -1 2 197632 ultralytics.nn.modules.block.C2f [128, 128, 2, True] \n",
|
|
" 7 -1 1 295424 ultralytics.nn.modules.conv.Conv [128, 256, 3, 2] \n",
|
|
" 8 -1 1 460288 ultralytics.nn.modules.block.C2f [256, 256, 1, True] \n",
|
|
" 9 -1 1 164608 ultralytics.nn.modules.block.SPPF [256, 256, 5] \n",
|
|
" 10 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest'] \n",
|
|
" 11 [-1, 6] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
|
" 12 -1 1 148224 ultralytics.nn.modules.block.C2f [384, 128, 1] \n",
|
|
" 13 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest'] \n",
|
|
" 14 [-1, 4] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
|
" 15 -1 1 37248 ultralytics.nn.modules.block.C2f [192, 64, 1] \n",
|
|
" 16 -1 1 36992 ultralytics.nn.modules.conv.Conv [64, 64, 3, 2] \n",
|
|
" 17 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
|
" 18 -1 1 123648 ultralytics.nn.modules.block.C2f [192, 128, 1] \n",
|
|
" 19 -1 1 147712 ultralytics.nn.modules.conv.Conv [128, 128, 3, 2] \n",
|
|
" 20 [-1, 9] 1 0 ultralytics.nn.modules.conv.Concat [1] \n",
|
|
" 21 -1 1 493056 ultralytics.nn.modules.block.C2f [384, 256, 1] \n",
|
|
" 22 [15, 18, 21] 1 752287 ultralytics.nn.modules.head.Detect [5, 16, None, [64, 128, 256]] \n",
|
|
"Model summary: 130 layers, 3,011,823 parameters, 3,011,807 gradients, 8.2 GFLOPs\n",
|
|
"\n",
|
|
"Transferred 319/355 items from pretrained weights\n",
|
|
"Freezing layer 'model.22.dfl.conv.weight'\n",
|
|
"\u001b[34m\u001b[1mtrain: \u001b[0mFast image access ✅ (ping: 0.0±0.0 ms, read: 1663.3±1491.0 MB/s, size: 29.1 KB)\n",
|
|
"\u001b[K\u001b[34m\u001b[1mtrain: \u001b[0mScanning /home/nathan/HadesStarYOLO/dataset_v5/labels.cache... 165 images, 0 backgrounds, 0 corrupt: 100% ━━━━━━━━━━━━ 165/165 57.7Mit/s 0.0s\n",
|
|
"\u001b[34m\u001b[1mval: \u001b[0mFast image access ✅ (ping: 0.0±0.0 ms, read: 1137.3±892.6 MB/s, size: 122.5 KB)\n",
|
|
"\u001b[K\u001b[34m\u001b[1mval: \u001b[0mScanning /home/nathan/HadesStarYOLO/dataset_v5/labels.cache... 165 images, 0 backgrounds, 0 corrupt: 100% ━━━━━━━━━━━━ 165/165 6.4Mit/s 0.0s\n",
|
|
"\u001b[34m\u001b[1moptimizer:\u001b[0m 'optimizer=auto' found, ignoring 'lr0=0.01' and 'momentum=0.937' and determining best 'optimizer', 'lr0' and 'momentum' automatically... \n",
|
|
"\u001b[34m\u001b[1moptimizer:\u001b[0m AdamW(lr=0.001111, momentum=0.9) with parameter groups 57 weight(decay=0.0), 64 weight(decay=0.0005), 63 bias(decay=0.0)\n",
|
|
"Plotting labels to /home/nathan/HadesStarYOLO/runs/detect/train-22/labels.jpg... \n",
|
|
"Image sizes 640 train, 640 val\n",
|
|
"Using 8 dataloader workers\n",
|
|
"Logging results to \u001b[1m/home/nathan/HadesStarYOLO/runs/detect/train-22\u001b[0m\n",
|
|
"Starting training for 50 epochs...\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 1/50 1.08G 2.321 4.195 1.342 23 640: 100% ━━━━━━━━━━━━ 42/42 3.0it/s 13.9s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.5it/s 2.8s0.1s\n",
|
|
" all 165 1385 0.00152 0.0802 0.00995 0.00372\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 2/50 1.5G 2.381 3.269 1.243 22 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.4it/s 2.8s0.1s\n",
|
|
" all 165 1385 0.193 0.0698 0.0897 0.0345\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 3/50 1.5G 2.294 2.919 1.206 11 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.1it/s 3.0s0.1s\n",
|
|
" all 165 1385 0.145 0.265 0.122 0.0526\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 4/50 1.5G 2.236 2.837 1.239 5 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.1it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.236 0.265 0.179 0.0837\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 5/50 1.5G 2.097 2.685 1.182 42 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.2s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.2it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.258 0.314 0.196 0.0898\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 6/50 1.5G 2.172 2.574 1.236 13 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.3it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.183 0.287 0.155 0.0649\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 7/50 1.5G 2.086 2.504 1.192 14 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.2it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.273 0.293 0.224 0.106\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 8/50 1.5G 2.058 2.468 1.152 17 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.3it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.319 0.312 0.253 0.118\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 9/50 1.5G 2.044 2.458 1.174 6 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.2s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.1it/s 3.0s0.1s\n",
|
|
" all 165 1385 0.35 0.385 0.301 0.151\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 10/50 1.5G 2.025 2.307 1.142 49 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.2s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.2it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.339 0.384 0.334 0.158\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 11/50 1.5G 2.001 2.215 1.146 5 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.1it/s 3.0s0.1s\n",
|
|
" all 165 1385 0.363 0.379 0.352 0.171\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 12/50 1.56G 1.995 2.167 1.124 17 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.3it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.414 0.453 0.395 0.195\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 13/50 1.56G 1.971 2.122 1.129 14 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.1it/s 3.0s0.1s\n",
|
|
" all 165 1385 0.417 0.439 0.429 0.216\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 14/50 1.56G 1.946 2.055 1.146 4 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.2it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.47 0.479 0.473 0.236\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 15/50 1.56G 1.917 2.042 1.126 69 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.3it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.492 0.473 0.476 0.228\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 16/50 1.56G 1.961 2.012 1.112 14 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.1it/s 3.0s0.1s\n",
|
|
" all 165 1385 0.477 0.485 0.499 0.254\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 17/50 1.56G 1.886 1.972 1.086 13 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.1it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.459 0.506 0.475 0.241\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 18/50 1.56G 1.862 1.877 1.096 9 640: 100% ━━━━━━━━━━━━ 42/42 3.1it/s 13.4s0.2s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.0it/s 3.0s0.1s\n",
|
|
" all 165 1385 0.54 0.546 0.566 0.291\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 19/50 1.56G 1.835 1.865 1.112 21 640: 100% ━━━━━━━━━━━━ 42/42 3.2it/s 13.3s0.6s\n",
|
|
"\u001b[K Class Images Instances Box(P R mAP50 mAP50-95): 100% ━━━━━━━━━━━━ 21/21 7.2it/s 2.9s0.1s\n",
|
|
" all 165 1385 0.58 0.584 0.602 0.308\n",
|
|
"\n",
|
|
" Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size\n",
|
|
"\u001b[K 20/50 1.56G 1.864 1.849 1.139 61 640: 92% ━━━━━━━━━━━─ 39/42 3.1it/s 13.0s<1.0s"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# === STEP 2: TRAIN A NANO YOLO MODEL ===\n",
|
|
"print(\"\\n-> Initializing YOLO training...\")\n",
|
|
"# 'yolov8n.pt' is the Nano model.\n",
|
|
"model = YOLO(\"yolov8n.pt\")\n",
|
|
"\n",
|
|
"# Train the model on your 130 annotations\n",
|
|
"model.train(\n",
|
|
" data=yaml_path,\n",
|
|
" epochs=50, \n",
|
|
" imgsz=640,\n",
|
|
" batch=4,\n",
|
|
" patience=5,\n",
|
|
" amp=not is_1050, # Conditionaly enable this when not on a potato GPU\n",
|
|
" verbose=True)\n",
|
|
"print(\"\\n-> Training complete!\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "62af05d8-7cc7-43dd-8a86-b3fb446046a4",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Model info\n",
|
|
"\n",
|
|
"## Version 5\n",
|
|
"\n",
|
|
"The model 14 is fairly good\n",
|
|
"\n",
|
|
"Model 16 was trained until regression however it does not appear much better\n",
|
|
"\n",
|
|
"Added more data."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "138e1424-8ecd-4b76-a487-345bcb3b69f1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# model = YOLO(\"runs/detect/train-14/weights/best.pt\") # Much better than others. Uses V5 dataset and more confidently identifies objects."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "693c7e66-9c2f-427f-b2a0-311c4e5e8a38",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# === STEP 3: EXTRACT A RANDOM FRAME AND TEST DETECTION ===\n",
|
|
"print(\"\\n-> Running full-frame inference test...\")\n",
|
|
"\n",
|
|
"# --- CUSTOMIZATION CONFIGURATION ---\n",
|
|
"CONFIDENCE_THRESHOLD = 0.30 # Only show detections above this threshold\n",
|
|
"ALPHA = 0.4 # Transparency: 0.0 (completely invisible) to 1.0 (completely solid)\n",
|
|
"# -----------------------------------\n",
|
|
"\n",
|
|
"# Grab a random frame from your video file\n",
|
|
"cap = cv2.VideoCapture(VIDEO_PATH)\n",
|
|
"if not cap.isOpened():\n",
|
|
" print(f\"Error: Could not open video {VIDEO_PATH}\")\n",
|
|
"else:\n",
|
|
" total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
|
|
" random_idx = random.randint(0, total_frames - 1)\n",
|
|
" cap.set(cv2.CAP_PROP_POS_FRAMES, random_idx)\n",
|
|
" ret, frame = cap.read()\n",
|
|
" cap.release()\n",
|
|
"\n",
|
|
" if ret:\n",
|
|
" # Run inference using your specified confidence threshold\n",
|
|
" results = model.predict(source=frame, conf=CONFIDENCE_THRESHOLD, imgsz=640)\n",
|
|
"\n",
|
|
" # 1. Generate the fully solid annotated frame (boxes, labels, etc.)\n",
|
|
" solid_annotated = results[0].plot()\n",
|
|
"\n",
|
|
" # 2. Blend the original frame with the annotated frame to create transparency\n",
|
|
" # Formula: blended = (frame * (1 - ALPHA)) + (solid_annotated * ALPHA)\n",
|
|
" annotated_frame = cv2.addWeighted(solid_annotated, ALPHA, frame, 1 - ALPHA, 0)\n",
|
|
"\n",
|
|
" plt.figure(figsize=(24, 16))\n",
|
|
"\n",
|
|
" # Convert BGR (OpenCV format) to RGB (Matplotlib format)\n",
|
|
" plt.imshow(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))\n",
|
|
" plt.title(f\"YOLO Object Detection Test (Frame: {random_idx})\", fontsize=14)\n",
|
|
" plt.axis(\"off\")\n",
|
|
" plt.show()\n",
|
|
" plt.savefig(\"lastimage.png\")\n",
|
|
" else:\n",
|
|
" print(\"Error: Could not read a random frame from the video.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "2fd1b368-0363-4a7e-8035-566adbd8db19",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import cv2\n",
|
|
"import random\n",
|
|
"\n",
|
|
"# === CONFIGURATION FOR THE VIDEO OUTPUT ===\n",
|
|
"OUTPUT_PATH = \"yolo_inference_output.mp4\"\n",
|
|
"SECONDS_TO_RECORD = 5\n",
|
|
"CONF_THRESHOLD = 0.25 # <-- Confidence threshold for detections\n",
|
|
"ALPHA = 0.4 # <-- Transparency/opacity level (0.0 = completely transparent, 1.0 = solid)\n",
|
|
"\n",
|
|
"print(\"\\n-> Processing video and compiling to file...\")\n",
|
|
"\n",
|
|
"cap = cv2.VideoCapture(VIDEO_PATH)\n",
|
|
"\n",
|
|
"if not cap.isOpened():\n",
|
|
" print(f\"Error: Could not open source video {VIDEO_PATH}\")\n",
|
|
"else:\n",
|
|
" # 1. Gather original video properties\n",
|
|
" fps = cap.get(cv2.CAP_PROP_FPS)\n",
|
|
" if fps == 0: fps = 30 # Fallback\n",
|
|
" print(f\"FPS: {fps}\")\n",
|
|
" \n",
|
|
" frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
|
|
" frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
|
|
" frames_to_process = int(fps * SECONDS_TO_RECORD)\n",
|
|
" \n",
|
|
" # 2. Define the codec and create VideoWriter object\n",
|
|
" fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n",
|
|
" out = cv2.VideoWriter(OUTPUT_PATH, fourcc, fps, (frame_width, frame_height))\n",
|
|
" \n",
|
|
" print(f\"Processing {SECONDS_TO_RECORD} seconds ({frames_to_process} frames)...\")\n",
|
|
" print(f\"Output will be saved to: {OUTPUT_PATH}\")\n",
|
|
"\n",
|
|
" processed_count = 0\n",
|
|
"\n",
|
|
" # Ensure total_frames is defined in your environment (e.g., total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)))\n",
|
|
" random_idx = random.randint(0, max(0, total_frames - (frames_to_process + 1)))\n",
|
|
" cap.set(cv2.CAP_PROP_POS_FRAMES, random_idx)\n",
|
|
" \n",
|
|
" try:\n",
|
|
" for i in range(frames_to_process):\n",
|
|
" ret, frame = cap.read()\n",
|
|
" if not ret:\n",
|
|
" print(\"End of source video reached early.\")\n",
|
|
" break\n",
|
|
" \n",
|
|
" # Run YOLO inference\n",
|
|
" results = model.predict(source=frame, conf=CONF_THRESHOLD, imgsz=640, verbose=False)\n",
|
|
" \n",
|
|
" # Generate the fully solid annotated frame (boxes, labels, etc.)\n",
|
|
" solid_annotated = results[0].plot()\n",
|
|
" \n",
|
|
" # Blend the original frame with the solid annotated frame to create transparency\n",
|
|
" annotated_frame = cv2.addWeighted(solid_annotated, ALPHA, frame, 1 - ALPHA, 0)\n",
|
|
" \n",
|
|
" # Write the blended frame to the output file\n",
|
|
" out.write(annotated_frame)\n",
|
|
" \n",
|
|
" processed_count += 1\n",
|
|
" # Simple, non-intrusive text update every 10 frames\n",
|
|
" if processed_count % 10 == 0:\n",
|
|
" print(f\"Progress: {processed_count}/{frames_to_process} frames processed...\", end=\"\\r\")\n",
|
|
" \n",
|
|
" finally:\n",
|
|
" # 3. CRITICAL: Release everything when done to avoid file corruption\n",
|
|
" cap.release()\n",
|
|
" out.release()\n",
|
|
" print(f\"\\n-> Done! Successfully saved {processed_count} frames to {OUTPUT_PATH}\")"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3 (ipykernel)",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.10.12"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|