{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13746387,"sourceType":"datasetVersion","datasetId":8747012},{"sourceId":13816899,"sourceType":"datasetVersion","datasetId":8620533},{"sourceId":677607,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":513841,"modelId":528480}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip uninstall -y tensorflow\n!uv pip install --no-deps --system --no-index --find-links='/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/setup' 'connected-components-3d'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet')\n\nimport os\nimport gc\nimport cv2\nimport torch\nimport traceback\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nfrom tqdm import tqdm\nfrom PIL import Image\nfrom collections import defaultdict\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nfrom shutil import copyfile\nfrom scipy.signal import resample,savgol_filter, medfilt\nimport torchvision.transforms as T\n\nfrom stage0_model import Net as Stage0Net\nfrom stage0_common import *","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Explained baseline for the *PhysioNet - Digitization of ECG Images* competition\n\nThis notebook contributes three innovations to the competition:\n1. It shows how to find the markers in a scanned ECG (object detection in class `MarkerFinder`).\n2. It visualizes how to extract the time series from images of types 3 and 11 with a top-down plane sweep (function `convert_scanned_color()`).\n3. It shows how to extract the time series from grayscale images with a neural network (function `convert_scanned_grayscale()`).\n\n## Training data overview\n\nThere are 977 electrocardiograms (ids) in train, accounting for 84 GByte. Every electrocardiogram has 9 PNG files and one CSV file.\n\nThere are nine image types per ECG:\n- 0001 Original color ECG image generated by ECG-image-kit.\n- 0003 Image printed in color and scanned in color. → processed by `convert_scanned_color()`\n- 0004 Image printed in color and scanned in black and white. → processed by `convert_scanned_grayscale()`\n- 0005 Mobile photos of color printed images.\n- 0006 Mobile photos of ECGs on the screen of laptop.\n- 0009 Mobile photos of stained and soaked printed ECGs.\n- 0010 Mobile photos of printed ECGs with extensive damage.\n- 0011 Scans of printed ECG images with mold in color. → processed by `convert_scanned_color()`\n- 0012 Scans of printed ECG images with mold in black and white. → processed by `convert_scanned_grayscale()`\n\nThe sampling frequencies in train are 250, 256, 500, 512, 1000, 1025 per second.","metadata":{}},{"cell_type":"code","source":"train = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/train.csv')\ntest = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/test.csv')\ntest_dir = Path(\"/kaggle/input/physionet-ecg-image-digitization/test\")\nsubmission = pd.read_parquet(\"/kaggle/input/physionet-ecg-image-digitization/sample_submission.parquet\")\nlabel_dict = {}\nfor idx, row in tqdm(train.iterrows(), total=len(train)):\n    label_dict[idx] = pd.read_csv(f'/kaggle/input/physionet-ecg-image-digitization/train/{row.id}/{row.id}.csv')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test['id'] = test['id'].astype(str) \ntest_id = test['id'].unique().tolist()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Train Size Images**","metadata":{}},{"cell_type":"code","source":"def get_unique_sizes(directory):\n    size_counts = defaultdict(int)\n    for root, _, files in os.walk(directory):\n        for file in files:\n            if file.lower().endswith(('.png', '.jpg', '.jpeg', 'JPG')):\n                try:\n                    with Image.open(os.path.join(root, file)) as img:\n                        size = img.size\n                        size_counts[size] += 1\n                except Exception as e:\n                    print(f\"Error {file}: {e}\")\n\n    return size_counts\n\nfolders = [\n    \"/kaggle/input/physionet-ecg-image-digitization/train\"\n]\n\nfor folder in folders:\n    print(f\"\\n📂 Folder: {folder}\")\n    sizes = get_unique_sizes(folder)\n\n    if not sizes:\n        print(\"No images or mistake in code\")\n        continue\n    \n    sorted_sizes = sorted(sizes.items(), key=lambda x: x[1], reverse=True)\n\n    print(\"┌───────────────┬───────────────┬─────────┐\")\n    print(\"│ Width (px)  │ Height (px) │ Quantity │\")\n    print(\"├───────────────┼───────────────┼─────────┤\")\n    for (w, h), count in sorted_sizes:\n        print(f\"│ {w:<13} │ {h:<13} │ {count:<7} │\")\n    print(\"└───────────────┴───────────────┴─────────┘\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Test Size Images**","metadata":{}},{"cell_type":"code","source":"folders = [\n    \"/kaggle/input/physionet-ecg-image-digitization/test\"\n]\n\nfor folder in folders:\n    print(f\"\\n📂 Folder: {folder}\")\n    sizes = get_unique_sizes(folder)\n\n    if not sizes:\n        print(\"No images or mistake in code\")\n        continue\n    \n    sorted_sizes = sorted(sizes.items(), key=lambda x: x[1], reverse=True)\n\n    print(\"┌───────────────┬───────────────┬─────────┐\")\n    print(\"│  Width (px)   │  Height (px)  │ Quantity│\")\n    print(\"├───────────────┼───────────────┼─────────┤\")\n    for (w, h), count in sorted_sizes:\n        print(f\"│ {w:<13} │ {h:<13} │ {count:<7} │\")\n    print(\"└───────────────┴───────────────┴─────────┘\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"idx = 0\nprint(train.id[idx])\n\nTRAIN_DIR = '/kaggle/input/physionet-ecg-image-digitization/train/'\nname = str(train.id[idx])\ndf_with_id0 = TRAIN_DIR + name + '/' + name + '.csv'\n\ndf = pd.read_csv(df_with_id0)\ndf.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df.shape","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for col in df.columns:\n    print(f'Col: {col}; NaN`s: {df[col].isnull().sum()}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_metadata = train[train['id'] == 7663343]\n\n# Check if signal length matches recording duration\nfs = train_metadata['fs'].values[0]\nsig_len = train_metadata['sig_len'].values[0]\nduration = sig_len / fs\n\nprint(f\"Signal duration: {duration} seconds\")\nprint(f\"Sampling frequency: {fs} Hz\")\nprint(f\"Number of samples: {sig_len}\")\n\n# Compare with what we see on the images\ndef analyze_ecg_image(image_path):\n    \"\"\"Analyze ECG image to determine characteristics\"\"\"\n    img = plt.imread(image_path)\n    print(f\"\\nAnalysis of {os.path.basename(image_path)}:\")\n    print(f\"Image size: {img.shape}\")\n    \n    # Can add analysis of grid, time markers, etc.\n    return img\n\n# Analyze the first image\nanalyze_ecg_image(TRAIN_DIR + '7663343/7663343-0001.png')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_image_type(filename):\n    \"\"\"Determine image type based on filename\"\"\"\n    type_mapping = {\n        '0001': 'original_color',\n        '0003': 'printed_scanned_color', \n        '0004': 'printed_scanned_bw',\n        '0005': 'mobile_photo_color',\n        '0006': 'mobile_photo_screen',\n        '0009': 'stained_soaked',\n        '0010': 'extensive_damage',\n        '0011': 'mold_color',\n        '0012': 'mold_bw'\n    }\n    \n    image_id = filename.split('-')[1].split('.')[0]\n    return type_mapping.get(image_id, 'unknown')\n\ndef has_artifacts(filename):\n    \"\"\"Determine if the image has artifacts\"\"\"\n    artifact_types = ['0009', '0010', '0011', '0012']\n    image_id = filename.split('-')[1].split('.')[0]\n    return image_id in artifact_types\n\n# Compare the original signal with different image versions\nfig, axes = plt.subplots(3, 3, figsize=(18, 12))\n\n# Plot original signal\ntime = np.arange(len(df['II'])) / train_metadata['fs'].values[0]\naxes[0,0].plot(time, df['II'], 'b-', linewidth=0.8)\naxes[0,0].set_title('Original ECG Signal (Lead II)')\naxes[0,0].set_xlabel('Time (s)')\naxes[0,0].set_ylabel('mV')\naxes[0,0].grid(True)\n\n# Display different image versions\nimage_files = [f for f in os.listdir(TRAIN_DIR + '7663343/') if f.endswith('.png')]\nfor i, img_file in enumerate(image_files[:8]):\n    row = (i + 1) // 3\n    col = (i + 1) % 3\n    \n    img_path = TRAIN_DIR + '7663343/' + img_file\n    img = plt.imread(img_path)\n    \n    axes[row, col].imshow(img)\n    axes[row, col].set_title(f'{get_image_type(img_file)}\\n{img_file}')\n    axes[row, col].axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Creating a DataFrame for compliance analysis\nimage_analysis = []\n\nfor img_file in sorted(image_files):\n    img_path = TRAIN_DIR + '7663343/' + img_file\n    img = plt.imread(img_path)\n    \n    image_analysis.append({\n        'image_file': img_file,\n        'image_id': img_file.split('-')[1].split('.')[0],\n        'image_shape': img.shape,\n        'image_type': get_image_type(img_file),\n        'has_artifacts': has_artifacts(img_file)\n    })\n\nimage_df = pd.DataFrame(image_analysis)\nimage_df.head(10)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Looping through the training dataset\n\nWe define a generator function `train_images_and_labels`, which loops through a subset of the training images.","metadata":{}},{"cell_type":"code","source":"def train_images_and_labels(start=None, end=None, image_types=None, use_tqdm=True):\n    \"\"\"Generator function which yields a subset of the training images.\n    \n    Parameters\n    start: start index of slice\n    end: end index of slice\n    image_types: list of image types to select\n    \"\"\"\n    t = train.iloc[start:end]\n    iterable = t.iterrows()\n    if use_tqdm:\n        iterable = tqdm(iterable, total=len(t))\n    for idx, row in iterable:\n        png_paths = sorted(glob(f'/kaggle/input/physionet-ecg-image-digitization/train/{row.id}/{row.id}-*.png'))\n        labels = label_dict[idx]\n        for path in png_paths:\n            img_type = int(path[-8:-4])\n            if image_types is None or img_type in image_types:\n                ima = cv2.imread(path)\n\n                # The following lines document the possible shapes for every image type in train\n                # Test files may be different\n                shape = ima.shape\n                assert len(shape) == 3\n                assert (img_type == 1) <= (shape == (1700, 2200, 3)) # 200 pixels per inch on Letter paper\n                assert (img_type == 3) <= (shape[0] == 1652)\n                assert (img_type == 4) <= (shape[0] == 1652)\n                assert (img_type == 5) <= (shape in {(3024, 4032, 3), (1344, 1008, 3), (4032, 3024, 3)})\n                assert (img_type == 6) <= ((shape == (4000, 3000, 3) or (shape == (3000, 4000, 3))))\n                assert (img_type == 9) <= ((shape == (3024, 4032, 3) or (shape == (4032, 3024, 3))))\n                assert (img_type == 10) <= ((shape == (3024, 4032, 3) or (shape == (4032, 3024, 3))))\n                assert (img_type == 11) <= (shape[0] == 1652)\n                assert (img_type == 12) <= (shape[0] == 1652)\n\n                yield idx, ima, img_type, labels","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def analyze_image_quality(image_path):\n    \"\"\"Analyze image quality metrics\"\"\"\n    img = cv2.imread(image_path)\n    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n    \n    # Calculate metrics\n    brightness = np.mean(img_gray)\n    contrast = np.std(img_gray)\n    \n    # Calculate noise (using Laplacian variance)\n    laplacian_var = cv2.Laplacian(img_gray, cv2.CV_64F).var()\n    \n    return {\n        'brightness': brightness,\n        'contrast': contrast,\n        'sharpness': laplacian_var\n    }\n\n# Compare quality across different image types\nquality_metrics = []\nfor img_file in image_files:\n    img_path = TRAIN_DIR + '7663343/' + img_file\n    metrics = analyze_image_quality(img_path)\n    metrics['image_type'] = get_image_type(img_file)\n    metrics['filename'] = img_file\n    quality_metrics.append(metrics)\n\nquality_df = pd.DataFrame(quality_metrics)\nprint(quality_df.groupby('image_type').mean(numeric_only=True))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Analyze ECG waveform characteristics for Lead II\ndef analyze_ecg_waveform(signal, fs):\n    \"\"\"Extract basic ECG waveform features\"\"\"\n    from scipy.signal import find_peaks\n    \n    # Find R-peaks (simplified)\n    peaks, _ = find_peaks(signal, height=np.percentile(signal, 80), distance=fs*0.5)\n    \n    if len(peaks) > 1:\n        rr_intervals = np.diff(peaks) / fs  # in seconds\n        heart_rate = 60 / np.mean(rr_intervals)  # BPM\n        \n        return {\n            'heart_rate': heart_rate,\n            'num_beats': len(peaks),\n            'rr_std': np.std(rr_intervals),\n            'signal_mean': np.mean(signal),\n            'signal_std': np.std(signal)\n        }\n    \n    return None\n\n# Apply to our signal\necg_features = analyze_ecg_waveform(df['II'].values, fs)\nif ecg_features:\n    print(\"ECG Features:\")\n    for key, value in ecg_features.items():\n        print(f\"  {key}: {value:.2f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Finding lead endpoints with MarkerFinder\n\nBefore decoding an image, it's good to know the coordinates of the 17 lead endpoints in the ECG. The following cell defines the class `MarkerFinder`, which determines these points. 13 points are found by the pattern matching function `cv2.matchTemplate()`; the right endpoints of the four lines are inferred as linear combinations of the other vectors.","metadata":{}},{"cell_type":"code","source":"class MarkerFinder:\n    \"\"\"This class finds the 13 markers in scanned ecg images and guesses the 4 line ends.\"\"\"\n    # From https://www.kaggle.com/code/ambrosm/ecg-original-explained-baseline\n    \n    def __init__(self, show_templates=False):\n        # Derive the templates from type 1 images\n        # np.max keeps the gridlines and markers and removes the ecg lines\n        ima = np.max([\n            cv2.imread('/kaggle/input/physionet-ecg-image-digitization/train/4292118763/4292118763-0001.png'),\n            cv2.imread('/kaggle/input/physionet-ecg-image-digitization/train/4289880010/4289880010-0001.png'),\n            cv2.imread('/kaggle/input/physionet-ecg-image-digitization/train/4284351157/4284351157-0001.png'),\n        ], axis=0)\n\n        # Template points in global coordinates of type 1 images\n        absolute_points = np.zeros((17, 2), dtype=int)\n        for i in range(3):\n            absolute_points[5 * i] = np.array([707 + 284 * i, 118]) # y, x\n            for j in range(1, 5):\n                absolute_points[5 * i + j] = np.array([707 + 284 * i, 118 + 492 * j])\n        absolute_points[5 * 3] = np.array([1535, 118])\n        absolute_points[5 * 3 + 1] = np.array([1535, 118 + 492 * 4])\n\n        # Top left corner of template rectangle\n        template_positions = [None] * 17\n        for i in range(len(absolute_points)):\n            if absolute_points[i][1] < 118 + 492 * 4:\n                if i % 5 == 0:\n                    template_positions[i] = (absolute_points[i][0] - 87, absolute_points[i][1] - 50) # y, x\n                else:\n                    template_positions[i] = (absolute_points[i][0] - 37, absolute_points[i][1] - 13)\n\n        # Height and width of the templates\n        template_sizes = np.array([(105, 60)] * 17) # height, width\n\n        # Transform the points to relative coordinates (inside the template)\n        template_points = [np.array([absolute_points[i][0] - template_positions[i][0],\n                                     absolute_points[i][1] - template_positions[i][1]])\n                           if template_positions[i] is not None\n                           else None\n                           for i in range(len(absolute_points))]\n\n        # Save the template matrices\n        templates = [None] * 17\n        for i in range(len(template_positions)):\n            if template_points[i] is not None:\n                template = (ima[template_positions[i][0]:template_positions[i][0]+template_sizes[i][0],\n                            template_positions[i][1]:template_positions[i][1]+template_sizes[i][1]])\n                templates[i] = template\n\n        # Plot the template matrices\n        if show_templates:\n            _, axs = plt.subplots(4, 4, figsize=(5, 7))\n            for i in range(len(template_positions)):\n                if template_points[i] is not None:\n                    template = templates[i].copy()\n                    cv2.rectangle(template,\n                                  (template_points[i][1]-1, template_points[i][0]-1),\n                                  (template_points[i][1]+1, template_points[i][0]+1), \n                                  [255, 0, 0], 2)\n                    axs[i // 5, i % 5].imshow(template)\n            for i in range(13, len(axs.ravel())):\n                axs.ravel()[i].axis('off')\n            plt.tight_layout()\n            plt.suptitle('The templates for the 13 markers', y=1.01)\n            plt.show()\n\n        self._absolute_points = absolute_points\n        self._template_positions = template_positions\n        self._template_sizes = template_sizes\n        self._template_points = template_points\n        self._templates = templates\n        \n    def find_markers(self, ima, warn=False, plot=False, title=''):\n        \"\"\"Return 17 markers as list of size-2 integer arrays (row, column)\n\n        Parameters:\n        ima: array of shape (1652, height, 3)\n        \"\"\"\n        \n        if ima.shape[0] != 1652:\n            raise ValueError(\"Implemented only for scanned images (image types 3, 4, 11, 12)\")\n\n        markers = np.full((17, 2), -1)\n\n        # Find 13 template-based markers\n        for j in range(len(self._templates)):\n            if self._template_points[j] is not None:\n                t = self._template_positions[j][0]-100\n                l = max(self._template_positions[j][1]-100, 0)\n                search_range = (ima[t:self._template_positions[j][0]+100+self._template_sizes[j][0],\n                                l:self._template_positions[j][1]+250+self._template_sizes[j][0]])\n                res = cv2.matchTemplate(search_range, self._templates[j], cv2.TM_CCOEFF)\n                min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)\n    \n                top_left = max_loc\n                if warn and max_val < 3e7:\n                    bottom_right = (top_left[0] + self._templates[j].shape[1],\n                                    top_left[1] + self._templates[j].shape[0])\n                    print(j, top_left, max_val)\n                    search_range = search_range.copy()\n                    cv2.rectangle(search_range, top_left, bottom_right, 0, 2)\n                    plt.imshow(search_range)\n                    plt.show()\n                markers[j] = np.array((t + top_left[1] + self._template_points[j][0],\n                                       l + top_left[0] + self._template_points[j][1]))\n\n        # Guess the ends of the first three lines (can be outside the bounding box of the image)\n        for i in range(3):\n            m = markers[5 * i + 3] * 2 - markers[5 * i + 2]\n            markers[5 * i + 4] = m\n\n        # Guess the end of the fourth line (can be outside the bounding box of the image)\n        markers[16] = ((markers[14] * (284 + 260) - markers[9] * 260) / 284).astype(int)\n\n        if plot:\n            ima = ima.copy()\n            for m in markers:\n                if m is not None:\n                    cv2.rectangle(ima, (m[1]-40, m[0]-40), (m[1]+40, m[0]+40), (255, 0, 0), 2)\n            # plt.figure(figsize=(12, 8))\n            plt.imshow(ima)\n            plt.title(title)\n            plt.show()\n\n        return markers\n        \n    @staticmethod\n    def lead_info(lead):\n        \"\"\"Specify which markers mark the begin and the end of a lead.\"\"\"\n        begin, end = {\n            'I': (0, 1),\n            'II-subset': (5, 6),\n            'III': (10, 11),\n            'aVR': (1, 2),\n            'aVL': (6, 7),\n            'aVF': (11, 12),\n            'V1': (2, 3),\n            'V2': (7, 8),\n            'V3': (12, 13),\n            'V4': (3, 4),\n            'V5': (8, 9),\n            'V6': (13, 14),\n            'II': (15, 16),\n        }[lead]\n        return begin // 5, begin, end\n\n    def demo(self, ima, warn=False, title=''):\n        \"\"\"Plot the image with red markers\"\"\"\n        markers = self.find_markers(ima, warn, plot=True, title=title)\n\nmf = MarkerFinder(show_templates=False)\n\nima = cv2.imread('/kaggle/input/physionet-ecg-image-digitization/train/1026034238/1026034238-0011.png') # correct\nmf.demo(ima, warn=False, title='Scanned ECG with 17 line endpoints')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Segmentation Model** ","metadata":{}},{"cell_type":"code","source":"global_dict = {\n    \"stage0_dir\": \"/kaggle/working/stage0\",\n    \"stage1_dir\": \"/kaggle/working/stage1\",\n    \"stage2_dir\": \"/kaggle/working/stage2\",\n}\n\ndef change_color(image_rgb):\n    hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV)\n    h, s, v = cv2.split(hsv)\n    \n    v_denoised = cv2.fastNlMeansDenoising(v, h=6)\n    \n    std = np.std(v_denoised)\n    clip_limit = max(1.0, min(3.5, 2.0 + std / 25))\n    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))\n    v_enhanced = clahe.apply(v_denoised)\n    \n    hsv_enhanced = cv2.merge([h, s, v_enhanced])\n    return cv2.cvtColor(hsv_enhanced, cv2.COLOR_HSV2RGB)\n\nstage0_dir = Path(global_dict[\"stage0_dir\"])\nstage0_dir.mkdir(exist_ok=True)\n\nstage0_net = Stage0Net(pretrained=False)\nstage0_net = load_net(stage0_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage0-last.checkpoint.pth')\nstage0_net.to(\"cuda:0\")\nstage0_net.eval()\n\nfor n, sample_id in enumerate(tqdm(test_id)):\n    path = test_dir / f'{sample_id}.png'\n    output_path = stage0_dir / f'{sample_id}.png'\n    \n    image_original = cv2.imread(str(path), cv2.IMREAD_COLOR)\n    image_original = cv2.cvtColor(image_original, cv2.COLOR_BGR2RGB)\n    image_for_model = change_color(image_original)\n    \n    batch = image_to_batch(image_for_model)\n\n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage0_net(batch)\n        \n        rotated, keypoint = output_to_predict(image_original, batch, output)\n        normalised, _, _ = normalise_by_homography(rotated, keypoint)\n        \n        cv2.imwrite(str(output_path), cv2.cvtColor(normalised, cv2.COLOR_RGB2BGR))\n    except Exception as e:\n        traceback.print_exc()\n        copyfile(path, output_path)\n\n\nfrom stage1_model import Net as Stage1Net\nfrom stage1_common import *\n\nstage0_dir = Path(global_dict[\"stage0_dir\"])\nstage1_dir = Path(global_dict[\"stage1_dir\"])\nstage1_dir.mkdir(exist_ok=True)\n\nstage1_net = Stage1Net(pretrained=False)\nstage1_net = load_net(stage1_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage1-last.checkpoint.pth')\nstage1_net.to(\"cuda:0\")\n\nfor n, sample_id in enumerate(tqdm(test_id)):\n    path = stage0_dir / f'{sample_id}.png'\n    output_path = stage1_dir / f'{sample_id}.png'\n    image = cv2.imread(path, cv2.IMREAD_COLOR_RGB)\n    batch = {'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)}\n\n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage1_net(batch)\n        gridpoint_xy, _ = output_to_predict(image, batch, output)\n        rectified = rectify_image(image, gridpoint_xy)\n        cv2.imwrite(output_path, cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR))\n    except:\n        traceback.print_exc()\n        copyfile(path, output_path)\n\nimport torchvision.transforms as T\nfrom stage2_model import *\nfrom stage2_common import *\nfrom scipy.signal import savgol_filter, medfilt\n\n\nclass Net3(nn.Module):\n    \n    def __init__(self, pretrained=True):\n        super(Net3, self).__init__()\n        encoder_dim = [64, 128, 256, 512]\n        decoder_dim = [128, 64, 32, 16]\n\n        self.encoder = timm.create_model(\n            model_name='resnet34.a3_in1k', pretrained=pretrained, in_chans=3, num_classes=0, global_pool=''\n        )\n\n        self.decoder = MyCoordUnetDecoder(\n            in_channel=encoder_dim[-1],\n            skip_channel=encoder_dim[:-1][::-1] + [0],\n            out_channel=decoder_dim,\n            scale=[2, 2, 2, 2]\n        )\n        self.pixel = nn.Conv2d(decoder_dim[-1], 4, 1)\n\n    def forward(self, image):\n        encode = encode_with_resnet(self.encoder, image)\n        last, _ = self.decoder(feature=encode[-1], skip=encode[:-1][::-1] + [None])\n        pixel = self.pixel(last)\n        return pixel\n\nstage1_dir = Path(global_dict[\"stage1_dir\"])\nstage2_dir = Path(global_dict[\"stage2_dir\"])\nstage2_dir.mkdir(exist_ok=True)\n\nstage2_net = Net3(pretrained=False).to(\"cuda:0\")\nmodel_path = \"/kaggle/input/physio-seg-public/pytorch/net3_009_4200/1/iter_0004200.pt\"\nstage2_net.load_state_dict(torch.load(model_path))\nstage2_net.eval()\n\nx0, x1 = 0, 2176\ny0, y1 = 0, 1696\nzero_mv = [703.5, 987.5, 1271.5, 1531.5]\nmv_to_pixel = 78.52\nt0, t1 = 235, 4161\n\nresize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)\n\nfor n, sample_id in enumerate(tqdm(test_id)):\n    path = stage1_dir / f'{sample_id}.png'\n    output_path = stage2_dir / f'{sample_id}.npy'\n    image = cv2.imread(path, cv2.IMREAD_COLOR_RGB)\n    \n    length = test[(test['id']==sample_id) & (test['lead']=='II')].iloc[0].number_of_rows\n    \n    image = image[y0:y1, x0:x1] / 255\n    batch = resize(torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)).float().to(\"cuda:0\")\n    \n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage2_net(batch)\n        \n        pixel = torch.sigmoid(output).float().data.cpu().numpy()[0]\n        series_in_pixel = pixel_to_series(pixel[..., t0:t1], zero_mv, length)\n        series = (np.array(zero_mv).reshape(4, 1) - series_in_pixel) / mv_to_pixel\n        \n       \n        for i in range(series.shape[0]):\n            series[i] = savgol_filter(series[i], window_length=7, polyorder=2)\n\n        np.save(output_path, series)\n    except:\n        traceback.print_exc()\n        series = np.zeros((4, length)) \n        np.save(output_path, series)\n\ndef series_dict(series):\n    d = {}\n    for l in range(3):\n        lead_names = [\n            ['I',   'aVR', 'V1', 'V4'],\n            ['II',  'aVL', 'V2', 'V5'],\n            ['III', 'aVF', 'V3', 'V6'],\n        ][l]\n        split = np.array_split(series[l], 4)\n        for (k, s) in zip(lead_names, split):\n            d[k] = s\n    \n    d['II'] = series[3]\n    \n    return d\n\n\nstage2_dir = Path(global_dict[\"stage2_dir\"])\n\nres = []\ngb = test.groupby('id')\n\nfor i, (sample_id, df) in enumerate(tqdm(gb)):\n    series = np.load(stage2_dir / f'{sample_id}.npy')\n    d_series = series_dict(series)\n\n    for _, d in df.iterrows():\n        s = d_series.get(d.lead, np.zeros(d.number_of_rows))\n        \n        if len(s) != d.number_of_rows:\n            x_old = np.linspace(0, 1, len(s))\n            x_new = np.linspace(0, 1, d.number_of_rows)\n            s = np.interp(x_new, x_old, s)\n        \n        row_id = [f'{sample_id}_{x}_{d.lead}' for x in range(d.number_of_rows)]\n        res.append(pd.DataFrame({'id': row_id, 'value': s}))\n\n    if i % 100 == 0:\n        gc.collect()\n\nsubmission = pd.concat(res, axis=0, ignore_index=True)\nsubmission.to_csv('submission.csv', index=False)\nsubmission.head(30)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}