{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":13517934,"sourceType":"datasetVersion","datasetId":8582983}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# ECG Signal Detection with YOLO\n\nWe manually annotated 100 ECG images to train a YOLOv8 model for detecting 4 signal strips per image. The trained model weights are used for inference.","metadata":{}},{"cell_type":"code","source":"!pip install ultralytics --quiet\nimport cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom ultralytics import YOLO\nfrom pathlib import Path\nimport random\n\n# Load trained model\nmodel = YOLO('/kaggle/input/signal-detector/ecg_signal_detector2/weights/best.pt')\n\n# Select 4 random test images\ntest_images = list(Path('/kaggle/input/physionet-ecg-image-digitization/train').rglob('*.png'))[:1000]  # Sample from first folders\nrandom.seed(13)\nselected_images = random.sample(test_images, 4)\n\nprint(f\"Model loaded. Testing on {len(selected_images)} images.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-27T10:27:39.737948Z","iopub.execute_input":"2025-10-27T10:27:39.738256Z","iopub.status.idle":"2025-10-27T10:27:41.233552Z","shell.execute_reply.started":"2025-10-27T10:27:39.738235Z","shell.execute_reply":"2025-10-27T10:27:41.232732Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Raw YOLO Predictions (No Filtering)\n\nUsing default confidence threshold (0.25), the model sometimes:\n- Detects <4 signals (misses some)\n- Detects >4 signals (false positives)\n- Detects partial signals instead of full strips","metadata":{}},{"cell_type":"code","source":"fig, axes = plt.subplots(2, 2, figsize=(14, 10))\naxes = axes.flatten()\n\nfor idx, img_path in enumerate(selected_images):\n    # Load image\n    img = cv2.imread(str(img_path))\n    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    \n    # Run YOLO with default settings\n    results = model(str(img_path), conf=0.25, verbose=False)\n    boxes = results[0].boxes\n    \n    # Draw boxes\n    for box in boxes:\n        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()\n        conf = box.conf[0].cpu().numpy()\n        cv2.rectangle(img_rgb, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 3)\n        cv2.putText(img_rgb, f'{conf:.2f}', (int(x1), int(y1)-10), \n                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)\n    \n    axes[idx].imshow(img_rgb)\n    axes[idx].set_title(f'{Path(img_path).name}\\n{len(boxes)} detections (expected 4)', fontsize=10)\n    axes[idx].axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-27T10:27:41.23492Z","iopub.execute_input":"2025-10-27T10:27:41.23526Z","iopub.status.idle":"2025-10-27T10:27:49.318586Z","shell.execute_reply.started":"2025-10-27T10:27:41.235241Z","shell.execute_reply":"2025-10-27T10:27:49.317446Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Note:\nHere it may look like all signals are detected properly by the YOLO model, but since i found to be defetcs, i have added below post-training steps to the model.","metadata":{}},{"cell_type":"markdown","source":"## Smart Filtering Strategy\n\nBased on annotation statistics:\n- All images have exactly 4 signals\n- Signal boxes cover ~86% of image width (σ=5.9%)\n- Box widths within same image vary <0.5%\n\n**Filtering approach:**\n1. Lower confidence threshold to 0.15 (catch all signals)\n2. Filter boxes by width (≥65% of image width)\n3. Keep top 4 boxes by confidence\n4. Sort by Y-coordinate (top→bottom = Lead 1→4)","metadata":{}},{"cell_type":"code","source":"def smart_filter_boxes(boxes, img_width, img_height, min_width_ratio=0.65, target_count=4):\n    \"\"\"Filter and sort boxes to get exactly 4 signal strips.\"\"\"\n    if len(boxes) == 0:\n        return []\n    \n    # Extract box data\n    box_data = []\n    for box in boxes:\n        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()\n        conf = box.conf[0].cpu().numpy()\n        width = x2 - x1\n        y_center = (y1 + y2) / 2\n        \n        box_data.append({\n            'x1': int(x1), 'y1': int(y1), 'x2': int(x2), 'y2': int(y2),\n            'conf': float(conf), 'y_center': y_center,\n            'width_ratio': width / img_width\n        })\n    \n    # Filter by width\n    filtered = [b for b in box_data if b['width_ratio'] >= min_width_ratio]\n    \n    # Keep top N by confidence\n    if len(filtered) > target_count:\n        filtered = sorted(filtered, key=lambda b: b['conf'], reverse=True)[:target_count]\n    \n    # Sort by Y position (top to bottom)\n    filtered = sorted(filtered, key=lambda b: b['y_center'])\n    \n    return filtered","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-27T10:26:53.883557Z","iopub.execute_input":"2025-10-27T10:26:53.884024Z","iopub.status.idle":"2025-10-27T10:26:53.892406Z","shell.execute_reply.started":"2025-10-27T10:26:53.88399Z","shell.execute_reply":"2025-10-27T10:26:53.891752Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Filtered Predictions\n\nApplying smart filtering to ensure exactly 4 signals are detected.","metadata":{}},{"cell_type":"code","source":"fig, axes = plt.subplots(2, 2, figsize=(14, 10))\naxes = axes.flatten()\n\ncolors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]  # Different color per lead\n\nfor idx, img_path in enumerate(selected_images):\n    # Load image\n    img = cv2.imread(str(img_path))\n    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    img_h, img_w = img.shape[:2]\n    \n    # Run YOLO with lower threshold\n    results = model(str(img_path), conf=0.15, iou=0.5, verbose=False)\n    boxes = results[0].boxes\n    \n    # Apply smart filtering\n    filtered_boxes = smart_filter_boxes(boxes, img_w, img_h)\n    \n    # Draw filtered boxes\n    for i, box in enumerate(filtered_boxes):\n        x1, y1, x2, y2 = box['x1'], box['y1'], box['x2'], box['y2']\n        conf = box['conf']\n        color = colors[i % 4]\n        \n        cv2.rectangle(img_rgb, (x1, y1), (x2, y2), color, 3)\n        cv2.putText(img_rgb, f'Lead {i+1}: {conf:.2f}', (x1, y1-10),\n                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)\n    \n    title_color = 'green' if len(filtered_boxes) == 4 else 'red'\n    axes[idx].imshow(img_rgb)\n    axes[idx].set_title(f'{Path(img_path).name}\\n{len(filtered_boxes)}/4 signals detected',\n                       fontsize=10, color=title_color)\n    axes[idx].axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-27T10:26:56.904541Z","iopub.execute_input":"2025-10-27T10:26:56.90482Z","iopub.status.idle":"2025-10-27T10:26:59.905136Z","shell.execute_reply.started":"2025-10-27T10:26:56.904799Z","shell.execute_reply":"2025-10-27T10:26:59.904167Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Results\n\nSmart filtering achieves 100% success rate on validation set (20 images):\n- **mAP50**: 0.937\n- **mAP50-95**: 0.589\n- **Precision**: 0.886\n- **Recall**: 0.871\n\nThe filtered approach ensures exactly 4 signals are detected and correctly ordered (top→bottom).","metadata":{}},{"cell_type":"code","source":"# Example: Extract signals from a single image\nimg_path = selected_images[0]\nimg = cv2.imread(str(img_path))\nimg_h, img_w = img.shape[:2]\n\nresults = model(str(img_path), conf=0.15, verbose=False)\nboxes = results[0].boxes\nfiltered_boxes = smart_filter_boxes(boxes, img_w, img_h)\n\n# Crop and display each signal\nfig, axes = plt.subplots(1, 4, figsize=(16, 3))\nfor i, box in enumerate(filtered_boxes):\n    x1, y1, x2, y2 = box['x1'], box['y1'], box['x2'], box['y2']\n    signal_crop = img[y1:y2, x1:x2]\n    signal_rgb = cv2.cvtColor(signal_crop, cv2.COLOR_BGR2RGB)\n    \n    axes[i].imshow(signal_rgb)\n    axes[i].set_title(f'Signal {i+1}', fontsize=10)\n    axes[i].axis('off')\n\nplt.tight_layout()\nplt.show()\n\nprint(f\"Extracted {len(filtered_boxes)} signals from {Path(img_path).name}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-27T10:26:59.906284Z","iopub.execute_input":"2025-10-27T10:26:59.906584Z","iopub.status.idle":"2025-10-27T10:27:00.427531Z","shell.execute_reply.started":"2025-10-27T10:26:59.906565Z","shell.execute_reply":"2025-10-27T10:27:00.426738Z"}},"outputs":[],"execution_count":null}]}