refactor: remove aitiautils

app.py

@@ -1,13 +1,14 @@
-# from io import BytesIO
+import os
+import tempfile
+import traceback as tb
+
 import gradio as gr
 import pandas as pd
 import numpy as np
 import torch
-import os
-from aitiautils.model_loader import ModelLoader
-import tempfile
 import matplotlib.pyplot as plt
-import traceback as tb
+
+from model import HRNetV2Wrapper
 
 # True 이면, tmp directory 에 파일 존재 유무와 상관없이 항상 새로운 이미지 생성
 ALWAYS_RECREATE_IMAGE = os.getenv("ALWAYS_RECREATE_IMAGE", "False").lower() == "true"
@@ -24,11 +25,14 @@ test_df = pd.read_csv("./res/ludb/dataset/test_for_public.csv").drop_duplicates(
 cutoffs = [0.001163482666015625, 0.15087890625, -0.587890625]
 lead_names = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
 
+hrnetv2_wrapper = HRNetV2Wrapper()
+
 
 def gen_seg(subject_id):
     input = np.load(f"./res/ludb/ecg_np/{subject_id}.npy")
-    network = ModelLoader("./res/models/hrnetv2/checkpoint.pth").get_network()
-    output: torch.Tensor = network(torch.from_numpy(input)).detach().numpy()
+    output: torch.Tensor = (
+        hrnetv2_wrapper.model(torch.from_numpy(input)).detach().numpy()
+    )
     seg = [(output[:, i, :] >= cutoffs[i]).astype(int) for i in range(len(cutoffs))]
     return input, np.stack(seg, axis=1)
 

model.py

@@ -0,0 +1,26 @@
+from res.impl.HRNetV2 import HRNetV2
+import torch
+
+
+class Config:
+    pass
+
+
+class HRNetV2Wrapper:
+    def __init__(self):
+        config = Config()
+        config.data_len = 5000
+        config.kernel_size = 5
+        config.dilation = 1
+        config.num_stages = 3
+        config.num_blocks = 6
+        config.num_modules = [1, 1, 1, 4, 3]
+        config.use_bottleneck = [1, 0, 0, 0, 0]
+        config.stage1_channels = 128
+        config.num_channels_init = 48
+        config.interpolate_mode = "linear"
+        config.output_size = 3
+        self.model = HRNetV2(config)
+        weights = torch.load("./res/models/hrnetv2/weights.pth")
+        self.model.load_state_dict(weights)
+        self.model = self.model.to("cpu").eval()

requirements.txt

@@ -1,8 +1,6 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
-
-# ./res/whl/aitiautils-1.1.59-py3-none-any.whl[torch]
-https://huggingface.co/spaces/MedicalAI-DP/ECG_Delineation/resolve/main/res/whl/aitiautils-1.1.59-py3-none-any.whl
 torch==2.0.0+cpu
 matplotlib==3.3.0
-netcal==1.3.5 
+pandas==2.2.3
+numpy==1.26.4
 gradio==5.12.0

res/models/hrnetv2/{checkpoint.pth → weights.pth}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2fd6e031d55efa48257225ebb2dce0e861fc02080830e655c61a618a227e1537
-size 44641506
+oid sha256:4d4d77219df74e4e69dc2feb33322caaf1b273353183e21c60385c9141d57bec
+size 22389961

res/whl/aitiautils-1.1.59-py3-none-any.whl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:99bb506c65d15b19d19691156efaa2fd15dcc78ea31913f817417c31d008100b
-size 14134456

utils.py

@@ -1,124 +0,0 @@
-def show_segmentation(ti=0, target_dir: str = "./seg"):
-    fig = plt.figure(figsize=(15, 2 * len(Alg) * 2))
-    fig.subplots_adjust(hspace=0, wspace=0.1)
-
-    # for lead_cnt, ti in enumerate(range(12 * target_idx, 12 * (target_idx + 1))):
-    test_df = pd.read_csv(test_path)
-    lead_type = ast.literal_eval(test_df.iloc[ti]["lead_type"])[0]
-    seq = test_df.iloc[ti]["seq"]
-    object_id: str = test_df.iloc[ti]["objectid"]
-    ecg_data_path = f"/bfai/data/ecg_data/{object_id[18:22]}/{object_id}.json"
-    # ecg_data_path = test_df.iloc[ti]["file_path"]
-    ecg_data = []
-    with open(ecg_data_path) as ecg_data_file:
-        ecg_json = json.load(ecg_data_file)
-        ecg_data = (
-            np.array(ecg_json["waveform"]["data"][lead_type])
-            * ecg_json["study"]["mv_unit"]
-        )
-        if seq != "1/1":
-            seq_idx, seq_range = [int(str_seq) for str_seq in seq.split("/")]
-            use_lead_length = 5000  # (500 * 10)
-            total_use_length = use_lead_length * seq_range
-            front_idx = int((len(ecg_data) - total_use_length) / 2)
-            ecg_data = ecg_data[front_idx:]
-
-    for alg, alg_idx in ALG_ORDER.items():
-        for pp_type_idx, pp_type in enumerate(["ori", "pp"]):
-            sub_fig = fig.add_subplot(len(Alg) * 2, 1, 2 * alg_idx + pp_type_idx + 1)
-            # sub_fig.set_title(f"{alg} - Lead {lead_type}")
-            sub_fig.text(
-                0.02,
-                0.5,
-                # f"{alg} {pp_type} - {lead_type}",
-                f"{alg} {pp_type}\n{lead_type}",
-                fontsize=9,
-                fontweight="bold",
-                ha="center",
-                va="center",
-                rotation=90,
-                transform=sub_fig.transAxes,
-            )
-            sub_fig.set_xticks([])
-            sub_fig.set_yticks([])
-
-            sub_fig.plot(range(len(ecg_data)), ecg_data, color="black", linewidth=1.0)
-            sub_fig.plot(
-                range(len(output[alg][ti][0])),
-                (output[alg][ti][0] >= cutoff[alg][0]).astype(int) / 2
-                if pp_type_idx == 0
-                else (pp_output[alg][ti][0]).astype(int) / 2,
-                label="P",
-                color="red",
-                linewidth=0.7,
-            )
-            sub_fig.plot(
-                range(len(output[alg][ti][1])),
-                (output[alg][ti][1] >= cutoff[alg][1]).astype(int)
-                if pp_type_idx == 0
-                else (pp_output[alg][ti][1]).astype(int),
-                label="QRS",
-                color="green",
-                linewidth=0.7,
-            )
-            sub_fig.plot(
-                (output[alg][ti][2] >= cutoff[alg][2]).astype(int) / 2
-                if pp_type_idx == 0
-                else (pp_output[alg][ti][2]).astype(int) / 2,
-                label="T",
-                color="blue",
-                linewidth=0.7,
-            )
-
-            sub_fig.plot(
-                range(len(origin_seg[ti][0])),
-                ((origin_seg[ti][0] > 0).astype(int) * (-1)) / 2,
-                label="P Label",
-                color="salmon",
-                linewidth=0.7,
-            )
-            sub_fig.plot(
-                range(len(origin_seg[ti][1])),
-                ((origin_seg[ti][1] > 0).astype(int) * (-1)),
-                label="QRS Label",
-                color="seagreen",
-                linewidth=0.7,
-            )
-            sub_fig.plot(
-                range(len(origin_seg[ti][2])),
-                ((origin_seg[ti][2] > 0).astype(int) * (-1)) / 2,
-                label="T Label",
-                color="darkslateblue",
-                linewidth=0.7,
-            )
-
-            sub_fig.plot(
-                range(len(origin_seg[ti][0])),
-                ((origin_seg[ti][0] < 0).astype(int) * (-1)),
-                label="P UnLabeled",
-                linestyle=":",
-                color="salmon",
-                linewidth=0.5,
-            )
-            sub_fig.plot(
-                range(len(origin_seg[ti][1])),
-                ((origin_seg[ti][1] < 0).astype(int) * (-1)),
-                label="QRS UnLabeled",
-                linestyle=":",
-                color="seagreen",
-                linewidth=0.5,
-            )
-            sub_fig.plot(
-                range(len(origin_seg[ti][2])),
-                ((origin_seg[ti][2] < 0).astype(int) * (-1)),
-                label="T UnLabeled",
-                linestyle=":",
-                color="darkslateblue",
-                linewidth=0.5,
-            )
-            # sub_fig.legend()
-    plt.savefig(
-        f"./{target_dir}/{test_df.iloc[ti]['lead_cnt']}_{object_id}_{lead_type}_{test_df.iloc[ti]['seq'].replace('/', '_')}.png",
-        dpi=150,
-    )
-    plt.close()