{ "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 }