Upload 30 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+CX0075_png.rf.f86acbac9d6c41151e8caed4914a3e89.jpg filter=lfs diff=lfs merge=lfs -text
+dev_3.jpg filter=lfs diff=lfs merge=lfs -text
+DP1983_png.rf.3f2a58f7f0feb4f9ad7b34149149553b.jpg filter=lfs diff=lfs merge=lfs -text
+HS1500_png.rf.8659b481c780f6b582532eb56d6f5349.jpg filter=lfs diff=lfs merge=lfs -text
+image_1.jpg filter=lfs diff=lfs merge=lfs -text
+image_3.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 yijingru
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,57 @@
----
-title: Cra Innovation Home App Demo
-emoji: 💻
-colorFrom: pink
-colorTo: pink
-sdk: streamlit
-sdk_version: 1.31.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Vertebra-Focused-Landmark-Detection-Pytorch
+Vertebra-Focused Landmark Detection for Scoliosis Assessment   [[arXiv](https://arxiv.org/pdf/2001.03187.pdf)]
+
+Accepted to ISBI2020.
+
+
+Please cite the article in your publications if it helps your research:
+
+	@article{yi2020vertebra,
+	  title={Vertebra-Focused Landmark Detection for Scoliosis Assessment},
+	  author={Yi, Jingru and Wu, Pengxiang and Huang, Qiaoying and Qu, Hui and Metaxas, Dimitris N},
+	  booktitle={ISBI},
+	  year={2020}
+	}
+
+
+<p align="center">
+	<img src="imgs/pic1.png", width="400">
+</p>
+
+<p align="center">
+	<img src="imgs/pic2.png", width="800">
+</p>
+
+# Dependencies
+Ubuntu 14.04, Python 3.6.4, PyTorch 1.1.0, OpenCV-Python 4.1.0.25 
+
+# How to start
+## Prepare Dataset
+To directly use dataset.py, you can arrange the dataset as follows:
+```
+/dataPath/data
+	/train/*.jpg
+	/val/*.jpg
+	/test/*.jpg
+/dataPath/labels/
+	/train/*.mat
+	/val/*.mat
+	/test/*.mat
+```
+The source dataset is from [[dataset16](http://spineweb.digitalimaginggroup.ca/spineweb/index.php?n=Main.Datasets#Dataset_16.3A_609_spinal_anterior-posterior_x-ray_images)].
+To adapt the code to your own dataset, you can modify the dataset.py, for example, change the 'load_gt_pts' function to adapt it to your own annotations. The pretrained weights can be downloaded [here](https://drive.google.com/drive/folders/1LhKnGVE8dUw0nK9_x4vPNY_L7sPY2_aQ?usp=sharing).
+	
+## Train the model
+```ruby
+python main.py --data_dir dataPath --epochs 50 --batch_size 2 --dataset spinal --phase train
+```
+
+## Test the model
+```ruby
+python main.py --resume weightPath --data_dir dataPath --dataset spinal  --phase test
+```
+
+
+## Evaluate the model
+```ruby
+python main.py --resume weightPath --data_dir dataPath --dataset spinal --phase eval

app.py

@@ -0,0 +1,399 @@
+import os
+import sys
+import shutil
+import importlib.util
+from io import BytesIO
+from ultralytics import YOLO
+from PIL import Image
+
+import torch
+# ─── FORCE CPU ONLY ─────────────────────────────────────────────────────────
+torch.Tensor.cuda      = lambda self, *args, **kwargs: self
+torch.nn.Module.cuda   = lambda self, *args, **kwargs: self
+torch.cuda.synchronize = lambda *args, **kwargs: None
+torch.cuda.is_available= lambda : False
+torch.cuda.device_count= lambda : 0
+_orig_to = torch.Tensor.to
+def _to_cpu(self, *args, **kwargs):
+    new_args = []
+    for a in args:
+        if isinstance(a, str) and a.lower().startswith("cuda"):
+            new_args.append("cpu")
+        elif isinstance(a, torch.device) and a.type=="cuda":
+            new_args.append(torch.device("cpu"))
+        else:
+            new_args.append(a)
+    if "device" in kwargs:
+        dev = kwargs["device"]
+        if (isinstance(dev, str) and dev.lower().startswith("cuda")) or \
+           (isinstance(dev, torch.device) and dev.type=="cuda"):
+            kwargs["device"] = torch.device("cpu")
+    return _orig_to(self, *new_args, **kwargs)
+torch.Tensor.to = _to_cpu
+
+from torch.utils.data import DataLoader as _DL
+def _dl0(ds, *a, **kw):
+    kw['num_workers'] = 0
+    return _DL(ds, *a, **kw)
+import torch.utils.data as _du
+_du.DataLoader = _dl0
+
+import cv2
+import numpy as np
+import streamlit as st
+from argparse import Namespace
+
+# ─── DYNAMIC IMPORT ─────────────────────────────────────────────────────────
+REPO = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(REPO)
+models_dir = os.path.join(REPO, "models")
+os.makedirs(models_dir, exist_ok=True)
+open(os.path.join(models_dir, "__init__.py"), "a").close()
+
+def load_mod(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    m    = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(m)
+    sys.modules[name] = m
+    return m
+
+dataset_mod = load_mod("dataset",     os.path.join(REPO, "dataset.py"))
+decoder_mod = load_mod("decoder",     os.path.join(REPO, "decoder.py"))
+draw_mod    = load_mod("draw_points", os.path.join(REPO, "draw_points.py"))
+test_mod    = load_mod("test",        os.path.join(REPO, "test.py"))
+load_mod("models.dec_net",     os.path.join(models_dir, "dec_net.py"))
+load_mod("models.model_parts", os.path.join(models_dir, "model_parts.py"))
+load_mod("models.resnet",      os.path.join(models_dir, "resnet.py"))
+load_mod("models.spinal_net",  os.path.join(models_dir, "spinal_net.py"))
+
+BaseDataset = dataset_mod.BaseDataset
+Network     = test_mod.Network
+
+# ─── STREAMLIT UI ───────────────────────────────────────────────────────────
+st.set_page_config(layout="wide", page_title="Vertebral Compression Fracture")
+
+st.markdown(
+        """
+    <div style='border: 2px solid #0080FF; border-radius: 5px; padding: 10px'>
+        <h1 style='text-align: center; color: #0080FF'>
+        🦴 Vertebral Compression Fracture Detection 🖼️
+        </h1>
+    </div>
+        """, unsafe_allow_html=True)
+st.markdown("")
+st.markdown("") 
+st.markdown("")
+col1, col2, col3, col4 = st.columns(4)
+
+with col4:
+    feature = st.selectbox(
+        "🔀 Select Feature",
+        ["How to use", "AP - Detection", "AP - Cobb angle" , "LA - Image Segmetation", "Contract"],
+        index=0,  # default to "AP"
+        help="Choose which view to display"
+    )
+
+if feature == "How to use":
+    st.markdown("## 📖 How to use this app")
+
+    col1, col2, col3 = st.columns(3)
+
+    with col1:
+        st.markdown(
+            """
+            <div style='border:2px solid #00BFFF; border-radius:10px; padding:15px; text-align:center; background-color:#F0F8FF'>
+                <h2>Step 1️⃣</h2>
+                <p>Go to <b>AP - Detection</b> or <b>LA - Image Segmentation</b></p>
+                <p>Select a sample image or upload your own image file.</p>
+                <p style='color:#008000;'><b>✅ Tip:</b> Best with X-ray images with clear vertebra visibility.</p>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
+
+    with col2:
+        st.markdown(
+            """
+            <div style='border:2px solid #00BFFF; border-radius:10px; padding:15px; text-align:center; background-color:#F0F8FF'>
+                <h2>Step 2️⃣</h2>
+                <p>Press the <b>Enter</b> button.</p>
+                <p>The system will process your image automatically.</p>
+                <p style='color:#FFA500;'><b>⏳ Note:</b> Processing time depends on image size.</p>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
+
+    with col3:
+        st.markdown(
+            """
+            <div style='border:2px solid #00BFFF; border-radius:10px; padding:15px; text-align:center; background-color:#F0F8FF'>
+                <h2>Step 3️⃣</h2>
+                <p>See the prediction results:</p>
+                <p style= text-align:left > 1. Bounding boxes & landmarks (AP)</p>
+                <p style= text-align:left >  2. Segmentation masks (LA)</p>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
+
+    st.markdown(" ")
+    st.info("สามารถเลือกฟีเจอร์ได้ผ่าน Select Feature โดยแต่ล่ะฟีเจอร์จะมีตัวอย่างกำกับให้ว่าเป็นยังไง")
+
+# store original dimensions
+elif feature == "AP - Detection":
+    uploaded = st.file_uploader("", type=["jpg", "jpeg", "png"])
+    orig_w = orig_h = None
+    img0 = None
+    run = st.button("Enter", use_container_width=True)
+    # ─── Maintain selected sample in session state ─────────
+    if "sample_img" not in st.session_state:
+        st.session_state.sample_img = None
+
+    # ─── SAMPLE BUTTONS ─────────────────────────────────────
+    with col1:
+        if st.button(" 1️⃣ Example",use_container_width=True):
+            st.session_state.sample_img = "image_1.jpg"
+    with col2:
+        if st.button(" 2️⃣ Example",use_container_width=True):
+            st.session_state.sample_img = "image_2.jpg"
+    with col3:
+        if st.button(" 3️⃣ Example",use_container_width=True):
+            st.session_state.sample_img = "image_3.jpg"
+
+    # ─── UI FOR UPLOAD + DISPLAY ───────────────────────────
+    col4, col5, col6 = st.columns(3)
+    with col4:
+        st.subheader("1️⃣ Upload & Run")
+
+        sample_img = st.session_state.sample_img  # read persisted choice
+
+        # case 1: uploaded file
+        if uploaded:
+            buf = uploaded.getvalue()
+            arr = np.frombuffer(buf, np.uint8)
+            img0 = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+            orig_h, orig_w = img0.shape[:2]
+            st.image(cv2.cvtColor(img0, cv2.COLOR_BGR2RGB), caption="Uploaded Image", use_container_width=True)
+
+        # case 2: selected sample image
+        elif sample_img is not None:
+            img_path = os.path.join(REPO, sample_img)
+            img0 = cv2.imread(img_path)
+            if img0 is not None:
+                orig_h, orig_w = img0.shape[:2]
+                st.image(cv2.cvtColor(img0, cv2.COLOR_BGR2RGB),
+                         caption=f"Sample Image: {sample_img}",
+                         use_container_width=True)
+            else:
+                st.error(f"Cannot find {sample_img} in directory!")
+
+
+
+    with col5:
+        st.subheader("2️⃣ Predictions")
+    with col6:
+        st.subheader("3️⃣ Heatmap")
+
+    # ─── ARGS & CHECKPOINT ─────────────────────────────────
+    args = Namespace(
+        resume="model_30.pth",
+        data_dir=os.path.join(REPO, "dataPath"),
+        dataset="spinal",
+        phase="test",
+        input_h=1024,
+        input_w=512,
+        down_ratio=4,
+        num_classes=1,
+        K=17,
+        conf_thresh=0.2,
+    )
+    weights_dir = os.path.join(REPO, "weights_spinal")
+    os.makedirs(weights_dir, exist_ok=True)
+    src_ckpt = os.path.join(REPO, "model_backup", args.resume)
+    dst_ckpt = os.path.join(weights_dir, args.resume)
+    if os.path.isfile(src_ckpt) and not os.path.isfile(dst_ckpt):
+        shutil.copy(src_ckpt, dst_ckpt)
+
+    # ─── MAIN LOGIC ────────────────────────────────────────
+    if img0 is not None and run and orig_w and orig_h:
+        # determine name for saving
+        if uploaded:
+            name = os.path.splitext(uploaded.name)[0] + ".jpg"
+        else:
+            name = os.path.splitext(sample_img)[0] + ".jpg"
+
+        testd = os.path.join(args.data_dir, "data", "test")
+        os.makedirs(testd, exist_ok=True)
+        cv2.imwrite(os.path.join(testd, name), img0)
+
+        orig_init = BaseDataset.__init__
+        def patched_init(self, data_dir, phase, input_h=None, input_w=None, down_ratio=4):
+            orig_init(self, data_dir, phase, input_h, input_w, down_ratio)
+            if phase == "test":
+                self.img_ids = [name]
+        BaseDataset.__init__ = patched_init
+
+        with st.spinner("Running model…"):
+            net = Network(args)
+            net.test(args, save=True)
+
+        out_dir = os.path.join(REPO, f"results_{args.dataset}")
+        pred_file = [f for f in os.listdir(out_dir)
+                     if f.startswith(name) and f.endswith("_pred.jpg")][0]
+        txtf = os.path.join(out_dir, f"{name}.txt")
+        imgf = os.path.join(out_dir, pred_file)
+
+        # ─── Annotated Predictions ─────────────────────────
+        base = cv2.imread(imgf)
+        txt = np.loadtxt(txtf)
+        tlx, tly = txt[:, 2].astype(int), txt[:, 3].astype(int)
+        trx, try_ = txt[:, 4].astype(int), txt[:, 5].astype(int)
+        blx, bly = txt[:, 6].astype(int), txt[:, 7].astype(int)
+        brx, bry = txt[:, 8].astype(int), txt[:, 9].astype(int)
+
+        top_pts, bot_pts, mids, dists = [], [], [], []
+        for (x1, y1), (x2, y2), (x3, y3), (x4, y4) in zip(
+                zip(tlx, tly), zip(trx, try_),
+                zip(blx, bly), zip(brx, bry)):
+            tm = np.array([(x1 + x2) / 2, (y1 + y2) / 2])
+            bm = np.array([(x3 + x4) / 2, (y3 + y4) / 2])
+            top_pts.append(tm)
+            bot_pts.append(bm)
+            mids.append((tm + bm) / 2)
+            dists.append(np.linalg.norm(bm - tm))
+
+        ref = dists[-1]
+        ann = base.copy()
+        for tm, bm in zip(top_pts, bot_pts):
+            cv2.line(ann, tuple(tm.astype(int)), tuple(bm.astype(int)), (0, 255, 255), 2)
+        for m, d in zip(mids, dists):
+            pct = (d - ref) / ref * 100
+            clr = (0, 255, 255) if pct <= 20 else (0, 165, 255) if pct <= 40 else (0, 0, 255)
+            pos = (int(m[0]) + 40, int(m[1]) + 5)
+            cv2.putText(ann, f"{pct:.0f}%", pos,
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, clr, 2, cv2.LINE_AA)
+
+        ann_resized = cv2.resize(ann, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
+        with col5:
+            st.image(cv2.cvtColor(ann_resized, cv2.COLOR_BGR2RGB), use_container_width=True)
+
+        H, W = base.shape[:2]
+        heat = np.zeros((H, W), np.float32)
+        for cx, cy in [(int(m[0]), int(m[1])) for m in mids]:
+            blob = np.zeros_like(heat)
+            blob[cy, cx] = 1.0
+            heat += cv2.GaussianBlur(blob, (0, 0), sigmaX=8, sigmaY=8)
+        heat /= (heat.max() + 1e-8)
+        hm8 = (heat * 255).astype(np.uint8)
+        hm_c = cv2.applyColorMap(hm8, cv2.COLORMAP_JET)
+
+        raw = cv2.imread(imgf, cv2.IMREAD_GRAYSCALE)
+        raw_b = cv2.cvtColor(raw, cv2.COLOR_GRAY2BGR)
+        overlay = cv2.addWeighted(raw_b, 0.6, hm_c, 0.4, 0)
+        overlay_resized = cv2.resize(overlay, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
+
+        with col6:
+            st.image(cv2.cvtColor(overlay_resized, cv2.COLOR_BGR2RGB), use_container_width=True)
+
+elif feature == "AP - Cobb angle":
+    st.write("กำลังพัฒนา")
+
+elif feature == "LA - Image Segmetation":
+    uploaded = st.file_uploader("", type=["jpg", "jpeg", "png"])
+    img0 = None
+
+    # ─── Maintain selected sample in session state ─────────
+    if "sample_img_la" not in st.session_state:
+        st.session_state.sample_img_la = None
+
+    # ─── SAMPLE BUTTONS ─────────────────────────────────────
+    with col1:
+        if st.button(" 1️⃣ Example ", use_container_width=True):
+            st.session_state.sample_img_la = "image_1_la.jpg"
+    with col2:
+        if st.button(" 2️⃣ Example ", use_container_width=True):
+            st.session_state.sample_img_la = "image_2_la.jpg"
+    with col3:
+        if st.button(" 3️⃣ Example ", use_container_width=True):
+            st.session_state.sample_img_la = "image_3_la.jpg"
+
+    # ─── UI FOR UPLOAD + DISPLAY ───────────────────────────
+    run_la = st.button("Enter", use_container_width=True)
+    col7, col8 = st.columns(2)
+
+    with col7:
+        st.subheader("🖼️ Original Image")
+
+        sample_img_la = st.session_state.sample_img_la  # read persisted choice
+
+        # case 1: uploaded file
+        if uploaded:
+            buf = uploaded.getvalue()
+            img0 = Image.open(BytesIO(buf)).convert("RGB")
+            st.image(img0, caption="Uploaded Image", use_container_width=True)
+
+        # case 2: selected sample image
+        elif sample_img_la is not None:
+            img_path = os.path.join(REPO, sample_img_la)
+            if os.path.isfile(img_path):
+                img0 = Image.open(img_path).convert("RGB")
+                st.image(img0, caption=f"Sample Image: {sample_img_la}", use_container_width=True)
+            else:
+                st.error(f"Cannot find {sample_img_la} in directory!")
+
+    with col8:
+        st.subheader("🔎 Predicted Image")
+
+        # ─── PREDICTION ────────────────────────────────────
+        if img0 is not None and run_la:
+            img_np = np.array(img0)
+            model = YOLO('./best.pt')  # or your correct path to best.pt
+            with st.spinner("Running YOLO model…"):
+                results = model(img_np, imgsz=640)
+            pred_img = results[0].plot(boxes=False, probs=False)  # returns numpy image with annotations
+            st.image(pred_img, caption="Prediction Result", use_container_width=True)
+
+elif feature == "Contract":
+    with col1:
+        st.image("dev_1.jpg", caption=None, use_container_width=True)
+        st.markdown(
+            """
+            <div style='border:2px solid #0080FF; border-radius:10px; padding:15px; text-align:center; background-color:#F0F8FF'>
+                <h3>Thitsanapat Uma</h3>
+                <a href='https://www.facebook.com/thitsanapat.uma' target='_blank'>
+                    🔗 Facebook Profile
+                </a>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
+    with col2:
+        st.image("dev_2.jpg", caption=None, use_container_width=True)
+        st.markdown(
+            """
+            <div style='border:2px solid #0080FF; border-radius:10px; padding:15px; text-align:center; background-color:#F0F8FF'>
+                <h3>Santipab Tongchan</h3>
+                <a href='https://www.facebook.com/santipab.tongchan.2025' target='_blank'>
+                    🔗 Facebook Profile
+                </a>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
+    with col3:
+        st.image("dev_3.jpg", caption=None, use_container_width=True)
+        st.markdown(
+            """
+            <div style='border:2px solid #0080FF; border-radius:10px; padding:15px; text-align:center; background-color:#F0F8FF'>
+                <h3>Suphanat Kamphapan</h3>
+                <a href='https://www.facebook.com/suphanat.kamphapan' target='_blank'>
+                    🔗 Facebook Profile
+                </a>
+            </div>
+            """,
+            unsafe_allow_html=True
+        )
+
+
+

best.pt

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

cobb_evaluate.py

@@ -0,0 +1,124 @@
+###########################################################################################
+## This code is transfered from matlab version of the MICCAI challenge
+## Oct 1 2019
+###########################################################################################
+import numpy as np
+import cv2
+
+
+def is_S(mid_p_v):
+    # mid_p_v:  34 x 2
+    ll = []
+    num = mid_p_v.shape[0]
+    for i in range(num-2):
+        term1 = (mid_p_v[i, 1]-mid_p_v[num-1, 1])/(mid_p_v[0, 1]-mid_p_v[num-1, 1])
+        term2 = (mid_p_v[i, 0]-mid_p_v[num-1, 0])/(mid_p_v[0, 0]-mid_p_v[num-1, 0])
+        ll.append(term1-term2)
+    ll = np.asarray(ll, np.float32)[:, np.newaxis]   # 32 x 1
+    ll_pair = np.matmul(ll, np.transpose(ll))        # 32 x 32
+    a = sum(sum(ll_pair))
+    b = sum(sum(abs(ll_pair)))
+    if abs(a-b)<1e-4:
+        return False
+    else:
+        return True
+
+def cobb_angle_calc(pts, image):
+    pts = np.asarray(pts, np.float32)   # 68 x 2
+    h,w,c = image.shape
+    num_pts = pts.shape[0]   # number of points, 68
+    vnum = num_pts//4-1
+
+    mid_p_v = (pts[0::2,:]+pts[1::2,:])/2   # 34 x 2
+    mid_p = []
+    for i in range(0, num_pts, 4):
+        pt1 = (pts[i,:]+pts[i+2,:])/2
+        pt2 = (pts[i+1,:]+pts[i+3,:])/2
+        mid_p.append(pt1)
+        mid_p.append(pt2)
+    mid_p = np.asarray(mid_p, np.float32)   # 34 x 2
+
+    for pt in mid_p:
+        cv2.circle(image,
+                   (int(pt[0]), int(pt[1])),
+                   12, (0,255,255), -1, 1)
+
+    for pt1, pt2 in zip(mid_p[0::2,:], mid_p[1::2,:]):
+        cv2.line(image,
+                 (int(pt1[0]), int(pt1[1])),
+                 (int(pt2[0]), int(pt2[1])),
+                 color=(0,0,255),
+                 thickness=5, lineType=1)
+
+    vec_m = mid_p[1::2,:]-mid_p[0::2,:]           # 17 x 2
+    dot_v = np.matmul(vec_m, np.transpose(vec_m)) # 17 x 17
+    mod_v = np.sqrt(np.sum(vec_m**2, axis=1))[:, np.newaxis]    # 17 x 1
+    mod_v = np.matmul(mod_v, np.transpose(mod_v)) # 17 x 17
+    cosine_angles = np.clip(dot_v/mod_v, a_min=0., a_max=1.)
+    angles = np.arccos(cosine_angles)   # 17 x 17
+    pos1 = np.argmax(angles, axis=1)
+    maxt = np.amax(angles, axis=1)
+    pos2 = np.argmax(maxt)
+    cobb_angle1 = np.amax(maxt)
+    cobb_angle1 = cobb_angle1/np.pi*180
+    flag_s = is_S(mid_p_v)
+    if not flag_s: # not S
+        # print('Not S')
+        cobb_angle2 = angles[0, pos2]/np.pi*180
+        cobb_angle3 = angles[vnum, pos1[pos2]]/np.pi*180
+        cv2.line(image,
+                 (int(mid_p[pos2 * 2, 0] ), int(mid_p[pos2 * 2, 1])),
+                 (int(mid_p[pos2 * 2 + 1, 0]), int(mid_p[pos2 * 2 + 1, 1])),
+                 color=(0, 255, 0), thickness=5, lineType=2)
+        cv2.line(image,
+                 (int(mid_p[pos1[pos2] * 2, 0]), int(mid_p[pos1[pos2] * 2, 1])),
+                 (int(mid_p[pos1[pos2] * 2 + 1, 0]), int(mid_p[pos1[pos2] * 2 + 1, 1])),
+                 color=(0, 255, 0), thickness=5, lineType=2)
+
+    else:
+        if (mid_p_v[pos2*2, 1]+mid_p_v[pos1[pos2]*2,1])<h:
+            # print('Is S: condition1')
+            angle2 = angles[pos2,:(pos2+1)]
+            cobb_angle2 = np.max(angle2)
+            pos1_1 = np.argmax(angle2)
+            cobb_angle2 = cobb_angle2/np.pi*180
+
+            angle3 = angles[pos1[pos2], pos1[pos2]:(vnum+1)]
+            cobb_angle3 = np.max(angle3)
+            pos1_2 = np.argmax(angle3)
+            cobb_angle3 = cobb_angle3/np.pi*180
+            pos1_2 = pos1_2 + pos1[pos2]-1
+
+            cv2.line(image,
+                     (int(mid_p[pos1_1 * 2, 0]), int(mid_p[pos1_1 * 2, 1])),
+                     (int(mid_p[pos1_1 * 2+1, 0]), int(mid_p[pos1_1 * 2 + 1, 1])),
+                     color=(0, 255, 0), thickness=5, lineType=2)
+
+            cv2.line(image,
+                     (int(mid_p[pos1_2 * 2, 0]), int(mid_p[pos1_2 * 2, 1])),
+                     (int(mid_p[pos1_2 * 2+1, 0]), int(mid_p[pos1_2 * 2 + 1, 1])),
+                     color=(0, 255, 0), thickness=5, lineType=2)
+
+        else:
+            # print('Is S: condition2')
+            angle2 = angles[pos2,:(pos2+1)]
+            cobb_angle2 = np.max(angle2)
+            pos1_1 = np.argmax(angle2)
+            cobb_angle2 = cobb_angle2/np.pi*180
+
+            angle3 = angles[pos1_1, :(pos1_1+1)]
+            cobb_angle3 = np.max(angle3)
+            pos1_2 = np.argmax(angle3)
+            cobb_angle3 = cobb_angle3/np.pi*180
+
+            cv2.line(image,
+                     (int(mid_p[pos1_1 * 2, 0]), int(mid_p[pos1_1 * 2, 1])),
+                     (int(mid_p[pos1_1 * 2+1, 0]), int(mid_p[pos1_1 * 2 + 1, 1])),
+                     color=(0, 255, 0), thickness=5, lineType=2)
+
+            cv2.line(image,
+                     (int(mid_p[pos1_2 * 2, 0]), int(mid_p[pos1_2 * 2, 1])),
+                     (int(mid_p[pos1_2 * 2+1, 0]), int(mid_p[pos1_2 * 2 + 1, 1])),
+                     color=(0, 255, 0), thickness=5, lineType=2)
+
+    return [cobb_angle1, cobb_angle2, cobb_angle3]

dataset.py

@@ -0,0 +1,88 @@
+import os
+import torch.utils.data as data
+import pre_proc
+import cv2
+from scipy.io import loadmat
+import numpy as np
+
+
+def rearrange_pts(pts):
+    boxes = []
+    for k in range(0, len(pts), 4):
+        pts_4 = pts[k:k+4,:]
+        x_inds = np.argsort(pts_4[:, 0])
+        pt_l = np.asarray(pts_4[x_inds[:2], :])
+        pt_r = np.asarray(pts_4[x_inds[2:], :])
+        y_inds_l = np.argsort(pt_l[:,1])
+        y_inds_r = np.argsort(pt_r[:,1])
+        tl = pt_l[y_inds_l[0], :]
+        bl = pt_l[y_inds_l[1], :]
+        tr = pt_r[y_inds_r[0], :]
+        br = pt_r[y_inds_r[1], :]
+        # boxes.append([tl, tr, bl, br])
+        boxes.append(tl)
+        boxes.append(tr)
+        boxes.append(bl)
+        boxes.append(br)
+    return np.asarray(boxes, np.float32)
+
+
+class BaseDataset(data.Dataset):
+    def __init__(self, data_dir, phase, input_h=None, input_w=None, down_ratio=4):
+        super(BaseDataset, self).__init__()
+        self.data_dir = data_dir
+        self.phase = phase
+        self.input_h = input_h
+        self.input_w = input_w
+        self.down_ratio = down_ratio
+        self.class_name = ['__background__', 'cell']
+        self.num_classes = 68
+        self.img_dir = os.path.join(data_dir, 'data', self.phase)
+        self.img_ids = sorted(os.listdir(self.img_dir))
+
+    def load_image(self, index):
+        image = cv2.imread(os.path.join(self.img_dir, self.img_ids[index]))
+        return image
+
+    def load_gt_pts(self, annopath):
+        pts = loadmat(annopath)['p2']   # num x 2 (x,y)
+        pts = rearrange_pts(pts)
+        return pts
+
+    def load_annoFolder(self, img_id):
+        return os.path.join(self.data_dir, 'labels', self.phase, img_id+'.mat')
+
+    def load_annotation(self, index):
+        img_id = self.img_ids[index]
+        annoFolder = self.load_annoFolder(img_id)
+        pts = self.load_gt_pts(annoFolder)
+        return pts
+
+    def __getitem__(self, index):
+        img_id = self.img_ids[index]
+        image = self.load_image(index)
+        if self.phase == 'test':
+            images = pre_proc.processing_test(image=image, input_h=self.input_h, input_w=self.input_w)
+            return {'images': images, 'img_id': img_id}
+        else:
+            aug_label = False
+            if self.phase == 'train':
+                aug_label = True
+            pts = self.load_annotation(index)   # num_obj x h x w
+            out_image, pts_2 = pre_proc.processing_train(image=image,
+                                                         pts=pts,
+                                                         image_h=self.input_h,
+                                                         image_w=self.input_w,
+                                                         down_ratio=self.down_ratio,
+                                                         aug_label=aug_label,
+                                                         img_id=img_id)
+
+            data_dict = pre_proc.generate_ground_truth(image=out_image,
+                                                       pts_2=pts_2,
+                                                       image_h=self.input_h//self.down_ratio,
+                                                       image_w=self.input_w//self.down_ratio,
+                                                       img_id=img_id)
+            return data_dict
+
+    def __len__(self):
+        return len(self.img_ids)

decoder.py

@@ -0,0 +1,77 @@
+import torch.nn.functional as F
+import numpy as np
+import torch
+
+class DecDecoder(object):
+    def __init__(self, K, conf_thresh):
+        self.K = 17
+        self.conf_thresh = conf_thresh
+
+    def _topk(self, scores):
+        batch, cat, height, width = scores.size()
+
+        topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), self.K)
+
+        topk_inds = topk_inds % (height * width)
+        topk_ys = (topk_inds / width).int().float()
+        topk_xs = (topk_inds % width).int().float()
+
+        topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), self.K)
+        topk_inds = self._gather_feat( topk_inds.view(batch, -1, 1), topk_ind).view(batch, self.K)
+        topk_ys = self._gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, self.K)
+        topk_xs = self._gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, self.K)
+
+        return topk_score, topk_inds, topk_ys, topk_xs
+
+
+    def _nms(self, heat, kernel=3):
+        hmax = F.max_pool2d(heat, (kernel, kernel), stride=1, padding=(kernel - 1) // 2)
+        keep = (hmax == heat).float()
+        return heat * keep
+
+    def _gather_feat(self, feat, ind, mask=None):
+        dim = feat.size(2)
+        ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
+        feat = feat.gather(1, ind)
+        if mask is not None:
+            mask = mask.unsqueeze(2).expand_as(feat)
+            feat = feat[mask]
+            feat = feat.view(-1, dim)
+        return feat
+
+    def _tranpose_and_gather_feat(self, feat, ind):
+        feat = feat.permute(0, 2, 3, 1).contiguous()
+        feat = feat.view(feat.size(0), -1, feat.size(3))
+        feat = self._gather_feat(feat, ind)
+        return feat
+
+    def ctdet_decode(self, heat, wh, reg):
+        # output: num_obj x 7
+        # 7: cenx, ceny, w, h, angle, score, cls
+        batch, c, height, width = heat.size()
+        heat = self._nms(heat)   # [1, 1, 256, 128]
+        scores, inds, ys, xs = self._topk(heat)
+        scores = scores.view(batch, self.K, 1)
+        reg = self._tranpose_and_gather_feat(reg, inds)
+        reg = reg.view(batch, self.K, 2)
+        xs = xs.view(batch, self.K, 1) + reg[:, :, 0:1]
+        ys = ys.view(batch, self.K, 1) + reg[:, :, 1:2]
+        wh = self._tranpose_and_gather_feat(wh, inds)
+        wh = wh.view(batch, self.K, 2*4)
+
+        tl_x = xs - wh[:,:,0:1]
+        tl_y = ys - wh[:,:,1:2]
+        tr_x = xs - wh[:,:,2:3]
+        tr_y = ys - wh[:,:,3:4]
+        bl_x = xs - wh[:,:,4:5]
+        bl_y = ys - wh[:,:,5:6]
+        br_x = xs - wh[:,:,6:7]
+        br_y = ys - wh[:,:,7:8]
+
+        pts = torch.cat([xs, ys,
+                         tl_x,tl_y,
+                         tr_x,tr_y,
+                         bl_x,bl_y,
+                         br_x,br_y,
+                         scores], dim=2).squeeze(0)
+        return pts.data.cpu().numpy()

draw_gaussian.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+
+def gaussian_radius(det_size, min_overlap=0.7):
+    height, width = det_size
+
+    a1  = 1
+    b1  = (height + width)
+    c1  = width * height * (1 - min_overlap) / (1 + min_overlap)
+    sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
+    r1  = (b1 + sq1) / 2
+
+    a2  = 4
+    b2  = 2 * (height + width)
+    c2  = (1 - min_overlap) * width * height
+    sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
+    r2  = (b2 + sq2) / 2
+
+    a3  = 4 * min_overlap
+    b3  = -2 * min_overlap * (height + width)
+    c3  = (min_overlap - 1) * width * height
+    sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
+    r3  = (b3 + sq3) / 2
+    return min(r1, r2, r3)
+
+def gaussian2D(shape, sigma=1):
+    m, n = [(ss - 1.) / 2. for ss in shape]
+    y, x = np.ogrid[-m:m+1,-n:n+1]
+
+    h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
+    h[h < np.finfo(h.dtype).eps * h.max()] = 0
+    return h
+
+def draw_umich_gaussian(heatmap, center, radius, k=1):
+    diameter = 2 * radius + 1
+    gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)
+
+    x, y = int(center[0]), int(center[1])
+
+    height, width = heatmap.shape[0:2]
+
+    left, right = min(x, radius), min(width - x, radius + 1)
+    top, bottom = min(y, radius), min(height - y, radius + 1)
+
+    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
+    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
+    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0:  # TODO debug
+        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
+    return heatmap

draw_loss.py

@@ -0,0 +1,64 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+
+def load_data(filename):
+    pts = []
+    f = open(filename, "rb")
+    for line in f:
+        pts.append(float(line.strip()))
+    f.close()
+    return pts
+
+dataset = 'spinal'
+weights_path = 'weights_'+dataset
+
+###############################################
+# Load data
+train_pts = load_data(os.path.join(weights_path, 'train_loss.txt'))
+val_pts = load_data(os.path.join(weights_path, 'val_loss.txt'))
+
+def draw_loss():
+    x = np.linspace(0, len(train_pts), len(train_pts))
+    plt.plot(x,train_pts,'ro-',label='train')
+    plt.plot(x,val_pts,'bo-',label='val')
+    # plt.axis([0, 50, 9.25, 11])
+    plt.legend(loc='upper right')
+
+    plt.xlabel('Epochs')
+    plt.ylabel('Loss')
+
+    plt.show()
+
+
+def draw_loss_ap():
+    ap05_pts = load_data(os.path.join(weights_path, 'ap_05_list.txt'))
+    ap07_pts = load_data(os.path.join(weights_path, 'ap_07_list.txt'))
+
+    x = np.linspace(0,len(train_pts),len(train_pts))
+    x1 = np.linspace(0, len(train_pts), len(ap05_pts))
+
+    fig, ax1 = plt.subplots()
+
+    color = 'tab:red'
+    ax1.set_xlabel('Epochs')
+    ax1.set_ylabel('Loss', color=color)
+    ax1.plot(x, train_pts, 'ro-',label='train')
+    ax1.plot(x, val_pts, 'bo-',label='val')
+    ax1.tick_params(axis='y', labelcolor=color)
+    plt.legend(loc = 'lower right')
+    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis
+    color = 'tab:blue'
+    ax2.set_ylabel('AP', color=color)  # we already handled the x-label with ax1
+    ax2.plot(x1, ap05_pts, 'go-',label='AP@05')
+    ax2.plot(x1, ap07_pts, 'yo-', label='AP@07')
+    ax2.tick_params(axis='y', labelcolor=color)
+
+    fig.tight_layout()  # otherwise the right y-label is slightly clipped
+    plt.legend(loc = 'upper right')
+    plt.show()
+
+
+if __name__ == '__main__':
+    draw_loss()
+    # draw_loss_ap()

draw_points.py

@@ -0,0 +1,92 @@
+import cv2
+import numpy as np
+
+colors = [[0.76590096, 0.0266074, 0.9806378],
+           [0.54197179, 0.81682527, 0.95081629],
+           [0.0799733, 0.79737015, 0.15173816],
+           [0.93240442, 0.8993321, 0.09901344],
+           [0.73130136, 0.05366301, 0.98405681],
+           [0.01664966, 0.16387004, 0.94158259],
+           [0.54197179, 0.81682527, 0.45081629],
+           # [0.92074915, 0.09919099 ,0.97590748],
+           [0.83445145, 0.97921679, 0.12250426],
+           [0.7300924, 0.23253621, 0.29764521],
+           [0.3856775, 0.94859286, 0.9910683],  # 10
+           [0.45762137, 0.03766411, 0.98755338],
+           [0.99496697, 0.09113071, 0.83322314],
+           [0.96478873, 0.0233309, 0.13149931],
+           [0.33240442, 0.9993321 , 0.59901344],
+            # [0.77690519,0.81783954,0.56220024],
+           # [0.93240442, 0.8993321, 0.09901344],
+           [0.95815068, 0.88436046, 0.55782268],
+           [0.03728425, 0.0618827, 0.88641827],
+           [0.05281129, 0.89572238, 0.08913828],
+
+           ]
+
+
+
+def draw_landmarks_regress_test(pts0, ori_image_regress, ori_image_points):
+    for i, pt in enumerate(pts0):
+        # color = np.random.rand(3)
+        color = colors[i]
+        # print(i+1, color)
+        color_255 = (255 * color[0], 255 * color[1], 255 * color[2])
+        cv2.circle(ori_image_regress, (int(pt[0]), int(pt[1])), 6, color_255, -1, 1)
+        # cv2.circle(ori_image, (int(pt[2]), int(pt[3])), 5, color_255, -1,1)
+        # cv2.circle(ori_image, (int(pt[4]), int(pt[5])), 5, color_255, -1,1)
+        # cv2.circle(ori_image, (int(pt[6]), int(pt[7])), 5, color_255, -1,1)
+        # cv2.circle(ori_image, (int(pt[8]), int(pt[9])), 5, color_255, -1,1)
+        cv2.arrowedLine(ori_image_regress, (int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.arrowedLine(ori_image_regress, (int(pt[0]), int(pt[1])), (int(pt[4]), int(pt[5])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.arrowedLine(ori_image_regress, (int(pt[0]), int(pt[1])), (int(pt[6]), int(pt[7])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.arrowedLine(ori_image_regress, (int(pt[0]), int(pt[1])), (int(pt[8]), int(pt[9])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.putText(ori_image_regress, '{}'.format(i + 1),
+                    (int(pt[4] + 10), int(pt[5] + 10)),
+                    cv2.FONT_HERSHEY_DUPLEX,
+                    1.2,
+                    color_255,  # (255,255,255),
+                    1,
+                    1)
+        # cv2.circle(ori_image, (int(pt[0]), int(pt[1])), 6, (255,255,255), -1,1)
+        cv2.circle(ori_image_points, (int(pt[2]), int(pt[3])), 5, color_255, -1, 1)
+        cv2.circle(ori_image_points, (int(pt[4]), int(pt[5])), 5, color_255, -1, 1)
+        cv2.circle(ori_image_points, (int(pt[6]), int(pt[7])), 5, color_255, -1, 1)
+        cv2.circle(ori_image_points, (int(pt[8]), int(pt[9])), 5, color_255, -1, 1)
+    return ori_image_regress, ori_image_points
+
+
+
+def draw_landmarks_pre_proc(out_image, pts):
+    for i in range(17):
+        pts_4 = pts[4 * i:4 * i + 4, :]
+        color = colors[i]
+        color_255 = (255 * color[0], 255 * color[1], 255 * color[2])
+        cv2.circle(out_image, (int(pts_4[0, 0]), int(pts_4[0, 1])), 5, color_255, -1, 1)
+        cv2.circle(out_image, (int(pts_4[1, 0]), int(pts_4[1, 1])), 5, color_255, -1, 1)
+        cv2.circle(out_image, (int(pts_4[2, 0]), int(pts_4[2, 1])), 5, color_255, -1, 1)
+        cv2.circle(out_image, (int(pts_4[3, 0]), int(pts_4[3, 1])), 5, color_255, -1, 1)
+    return np.uint8(out_image)
+
+
+def draw_regress_pre_proc(out_image, pts):
+    for i in range(17):
+        pts_4 = pts[4 * i:4 * i + 4, :]
+        pt = np.mean(pts_4, axis=0)
+        color = colors[i]
+        color_255 = (255 * color[0], 255 * color[1], 255 * color[2])
+        cv2.arrowedLine(out_image, (int(pt[0]), int(pt[1])), (int(pts_4[0, 0]), int(pts_4[0, 1])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.arrowedLine(out_image, (int(pt[0]), int(pt[1])), (int(pts_4[1, 0]), int(pts_4[1, 1])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.arrowedLine(out_image, (int(pt[0]), int(pt[1])), (int(pts_4[2, 0]), int(pts_4[2, 1])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.arrowedLine(out_image, (int(pt[0]), int(pt[1])), (int(pts_4[3, 0]), int(pts_4[3, 1])), color_255, 2, 1,
+                        tipLength=0.2)
+        cv2.putText(out_image, '{}'.format(i + 1), (int(pts_4[1, 0] + 10), int(pts_4[1, 1] + 10)),
+                    cv2.FONT_HERSHEY_DUPLEX, 1.2, color_255, 1, 1)
+    return np.uint8(out_image)

eval.py

@@ -0,0 +1,222 @@
+import torch
+import numpy as np
+from models import spinal_net
+import decoder
+import os
+from dataset import BaseDataset
+import time
+import cobb_evaluate
+
+def apply_mask(image, mask, alpha=0.5):
+    """Apply the given mask to the image.
+    """
+    color = np.random.rand(3)
+    for c in range(3):
+        image[:, :, c] = np.where(mask == 1,
+                                  image[:, :, c] *
+                                  (1 - alpha) + alpha * color[c] * 255,
+                                  image[:, :, c])
+    return image
+
+class Network(object):
+    def __init__(self, args):
+        torch.manual_seed(317)
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        heads = {'hm': args.num_classes,  # cen, tl, tr, bl, br
+                 'reg': 2*args.num_classes,
+                 'wh': 2*4,}
+
+        self.model = spinal_net.SpineNet(heads=heads,
+                                         pretrained=True,
+                                         down_ratio=args.down_ratio,
+                                         final_kernel=1,
+                                         head_conv=256)
+        self.num_classes = args.num_classes
+        self.decoder = decoder.DecDecoder(K=args.K, conf_thresh=args.conf_thresh)
+        self.dataset = {'spinal': BaseDataset}
+
+    def load_model(self, model, resume):
+        checkpoint = torch.load(resume, map_location=lambda storage, loc: storage)
+        print('loaded weights from {}, epoch {}'.format(resume, checkpoint['epoch']))
+        state_dict_ = checkpoint['state_dict']
+        model.load_state_dict(state_dict_, strict=False)
+        return model
+
+
+    def eval(self, args, save):
+        save_path = 'weights_'+args.dataset
+        self.model = self.load_model(self.model, os.path.join(save_path, args.resume))
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        dataset_module = self.dataset[args.dataset]
+        dsets = dataset_module(data_dir=args.data_dir,
+                               phase='test',
+                               input_h=args.input_h,
+                               input_w=args.input_w,
+                               down_ratio=args.down_ratio)
+
+        data_loader = torch.utils.data.DataLoader(dsets,
+                                                  batch_size=1,
+                                                  shuffle=False,
+                                                  num_workers=1,
+                                                  pin_memory=True)
+
+        total_time = []
+        landmark_dist = []
+        pr_cobb_angles = []
+        gt_cobb_angles = []
+        for cnt, data_dict in enumerate(data_loader):
+            begin_time = time.time()
+            images = data_dict['images'][0]
+            img_id = data_dict['img_id'][0]
+            images = images.to('cuda')
+            print('processing {}/{} image ...'.format(cnt, len(data_loader)))
+
+            with torch.no_grad():
+                output = self.model(images)
+                hm = output['hm']
+                wh = output['wh']
+                reg = output['reg']
+            torch.cuda.synchronize(self.device)
+            pts2 = self.decoder.ctdet_decode(hm, wh, reg)   # 17, 11
+            pts0 = pts2.copy()
+            pts0[:,:10] *= args.down_ratio
+            x_index = range(0,10,2)
+            y_index = range(1,10,2)
+            ori_image = dsets.load_image(dsets.img_ids.index(img_id)).copy()
+            h,w,c = ori_image.shape
+            pts0[:, x_index] = pts0[:, x_index]/args.input_w*w
+            pts0[:, y_index] = pts0[:, y_index]/args.input_h*h
+            # sort the y axis
+            sort_ind = np.argsort(pts0[:,1])
+            pts0 = pts0[sort_ind]
+            pr_landmarks = []
+            for i, pt in enumerate(pts0):
+                pr_landmarks.append(pt[2:4])
+                pr_landmarks.append(pt[4:6])
+                pr_landmarks.append(pt[6:8])
+                pr_landmarks.append(pt[8:10])
+            pr_landmarks = np.asarray(pr_landmarks, np.float32)   #[68, 2]
+
+            end_time = time.time()
+            total_time.append(end_time-begin_time)
+
+            gt_landmarks = dsets.load_gt_pts(dsets.load_annoFolder(img_id))
+            for pr_pt, gt_pt in zip(pr_landmarks, gt_landmarks):
+                    landmark_dist.append(np.sqrt((pr_pt[0]-gt_pt[0])**2+(pr_pt[1]-gt_pt[1])**2))
+
+            pr_cobb_angles.append(cobb_evaluate.cobb_angle_calc(pr_landmarks, ori_image))
+            gt_cobb_angles.append(cobb_evaluate.cobb_angle_calc(gt_landmarks, ori_image))
+
+        pr_cobb_angles = np.asarray(pr_cobb_angles, np.float32)
+        gt_cobb_angles = np.asarray(gt_cobb_angles, np.float32)
+
+        out_abs = abs(gt_cobb_angles - pr_cobb_angles)
+        out_add = gt_cobb_angles + pr_cobb_angles
+
+        term1 = np.sum(out_abs, axis=1)
+        term2 = np.sum(out_add, axis=1)
+
+        SMAPE = np.mean(term1 / term2 * 100)
+
+        print('mse of landmarkds is {}'.format(np.mean(landmark_dist)))
+        print('SMAPE is {}'.format(SMAPE))
+
+        total_time = total_time[1:]
+        print('avg time is {}'.format(np.mean(total_time)))
+        print('FPS is {}'.format(1./np.mean(total_time)))
+
+
+    def SMAPE_single_angle(self, gt_cobb_angles, pr_cobb_angles):
+        out_abs = abs(gt_cobb_angles - pr_cobb_angles)
+        out_add = gt_cobb_angles + pr_cobb_angles
+
+        term1 = out_abs
+        term2 = out_add
+
+        term2[term2==0] += 1e-5
+
+        SMAPE = np.mean(term1 / term2 * 100)
+        return SMAPE
+
+    def eval_three_angles(self, args, save):
+        save_path = 'weights_'+args.dataset
+        self.model = self.load_model(self.model, os.path.join(save_path, args.resume))
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        dataset_module = self.dataset[args.dataset]
+        dsets = dataset_module(data_dir=args.data_dir,
+                               phase='test',
+                               input_h=args.input_h,
+                               input_w=args.input_w,
+                               down_ratio=args.down_ratio)
+
+        data_loader = torch.utils.data.DataLoader(dsets,
+                                                  batch_size=1,
+                                                  shuffle=False,
+                                                  num_workers=1,
+                                                  pin_memory=True)
+
+        total_time = []
+        landmark_dist = []
+        pr_cobb_angles = []
+        gt_cobb_angles = []
+        for cnt, data_dict in enumerate(data_loader):
+            begin_time = time.time()
+            images = data_dict['images'][0]
+            img_id = data_dict['img_id'][0]
+            images = images.to('cuda')
+            print('processing {}/{} image ...'.format(cnt, len(data_loader)))
+
+            with torch.no_grad():
+                output = self.model(images)
+                hm = output['hm']
+                wh = output['wh']
+                reg = output['reg']
+            torch.cuda.synchronize(self.device)
+            pts2 = self.decoder.ctdet_decode(hm, wh, reg)   # 17, 11
+            pts0 = pts2.copy()
+            pts0[:,:10] *= args.down_ratio
+            x_index = range(0,10,2)
+            y_index = range(1,10,2)
+            ori_image = dsets.load_image(dsets.img_ids.index(img_id)).copy()
+            h,w,c = ori_image.shape
+            pts0[:, x_index] = pts0[:, x_index]/args.input_w*w
+            pts0[:, y_index] = pts0[:, y_index]/args.input_h*h
+            # sort the y axis
+            sort_ind = np.argsort(pts0[:,1])
+            pts0 = pts0[sort_ind]
+            pr_landmarks = []
+            for i, pt in enumerate(pts0):
+                pr_landmarks.append(pt[2:4])
+                pr_landmarks.append(pt[4:6])
+                pr_landmarks.append(pt[6:8])
+                pr_landmarks.append(pt[8:10])
+            pr_landmarks = np.asarray(pr_landmarks, np.float32)   #[68, 2]
+
+            end_time = time.time()
+            total_time.append(end_time-begin_time)
+
+            gt_landmarks = dsets.load_gt_pts(dsets.load_annoFolder(img_id))
+            for pr_pt, gt_pt in zip(pr_landmarks, gt_landmarks):
+                    landmark_dist.append(np.sqrt((pr_pt[0]-gt_pt[0])**2+(pr_pt[1]-gt_pt[1])**2))
+
+            pr_cobb_angles.append(cobb_evaluate.cobb_angle_calc(pr_landmarks, ori_image))
+            gt_cobb_angles.append(cobb_evaluate.cobb_angle_calc(gt_landmarks, ori_image))
+
+        pr_cobb_angles = np.asarray(pr_cobb_angles, np.float32)
+        gt_cobb_angles = np.asarray(gt_cobb_angles, np.float32)
+
+
+        print('SMAPE1 is {}'.format(self.SMAPE_single_angle(gt_cobb_angles[:,0], pr_cobb_angles[:,0])))
+        print('SMAPE2 is {}'.format(self.SMAPE_single_angle(gt_cobb_angles[:,1], pr_cobb_angles[:,1])))
+        print('SMAPE3 is {}'.format(self.SMAPE_single_angle(gt_cobb_angles[:,2], pr_cobb_angles[:,2])))
+
+        print('mse of landmarkds is {}'.format(np.mean(landmark_dist)))
+
+        total_time = total_time[1:]
+        print('avg time is {}'.format(np.mean(total_time)))
+        print('FPS is {}'.format(1./np.mean(total_time)))
+

loss.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RegL1Loss(nn.Module):
+    def __init__(self):
+        super(RegL1Loss, self).__init__()
+
+    def _gather_feat(self, feat, ind, mask=None):
+        dim = feat.size(2)
+        ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
+        feat = feat.gather(1, ind)
+        if mask is not None:
+            mask = mask.unsqueeze(2).expand_as(feat)
+            feat = feat[mask]
+            feat = feat.view(-1, dim)
+        return feat
+
+    def _tranpose_and_gather_feat(self, feat, ind):
+        feat = feat.permute(0, 2, 3, 1).contiguous()
+        feat = feat.view(feat.size(0), -1, feat.size(3))
+        feat = self._gather_feat(feat, ind)
+        return feat
+
+    def forward(self, output, mask, ind, target):
+        pred = self._tranpose_and_gather_feat(output, ind)
+        mask = mask.unsqueeze(2).expand_as(pred).float()
+        loss = F.l1_loss(pred * mask, target * mask, reduction='sum')
+        loss = loss / (mask.sum() + 1e-4)
+        return loss
+
+class FocalLoss(nn.Module):
+  def __init__(self):
+    super(FocalLoss, self).__init__()
+
+  def forward(self, pred, gt):
+      pos_inds = gt.eq(1).float()
+      neg_inds = gt.lt(1).float()
+      neg_weights = torch.pow(1 - gt, 4)
+
+      loss = 0
+
+      pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
+      neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds
+
+      num_pos  = pos_inds.float().sum()
+      pos_loss = pos_loss.sum()
+      neg_loss = neg_loss.sum()
+
+      if num_pos == 0:
+        loss = loss - neg_loss
+      else:
+        loss = loss - (pos_loss + neg_loss) / num_pos
+      return loss
+
+class LossAll(torch.nn.Module):
+    def __init__(self):
+        super(LossAll, self).__init__()
+        self.L_hm = FocalLoss()
+        self.L_off = RegL1Loss()
+        self.L_wh =  RegL1Loss()
+
+    def forward(self, pr_decs, gt_batch):
+        hm_loss  = self.L_hm(pr_decs['hm'],  gt_batch['hm'])
+        wh_loss  = self.L_wh(pr_decs['wh'], gt_batch['reg_mask'], gt_batch['ind'], gt_batch['wh'])
+        off_loss = self.L_off(pr_decs['reg'], gt_batch['reg_mask'], gt_batch['ind'], gt_batch['reg'])
+        loss_dec = hm_loss + off_loss + wh_loss
+        return loss_dec

main.py

@@ -0,0 +1,40 @@
+import argparse
+import train
+import test
+import eval
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='CenterNet Modification Implementation')
+    parser.add_argument('--num_epoch', type=int, default=50, help='Number of epochs')
+    parser.add_argument('--batch_size', type=int, default=2, help='Number of epochs')
+    parser.add_argument('--num_workers', type=int, default=4, help='Number of workers')
+    parser.add_argument('--init_lr', type=float, default=1.25e-4, help='Init learning rate')
+    parser.add_argument('--down_ratio', type=int, default=4, help='down ratio')
+    parser.add_argument('--input_h', type=int, default=1024, help='input height')
+    parser.add_argument('--input_w', type=int, default=512, help='input width')
+    parser.add_argument('--K', type=int, default=100, help='maximum of objects')
+    parser.add_argument('--conf_thresh', type=float, default=0.2, help='confidence threshold')
+    parser.add_argument('--seg_thresh', type=float, default=0.5, help='confidence threshold')
+    parser.add_argument('--num_classes', type=int, default=1, help='number of classes')
+    parser.add_argument('--ngpus', type=int, default=0, help='number of gpus')
+    parser.add_argument('--resume', type=str, default='model_last.pth', help='weights to be resumed')
+    parser.add_argument('--data_dir', type=str, default='../../Datasets/spinal/', help='data directory')
+    parser.add_argument('--phase', type=str, default='test', help='data directory')
+    parser.add_argument('--dataset', type=str, default='spinal', help='data directory')
+    args = parser.parse_args()
+    return args
+
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    if args.phase == 'train':
+        is_object = train.Network(args)
+        is_object.train_network(args)
+    elif args.phase == 'test':
+        is_object = test.Network(args)
+        is_object.test(args, save=True)
+    elif args.phase == 'eval':
+        is_object = eval.Network(args)
+        is_object.eval(args, save=True)
+        # is_object.eval_three_angles(args, save=False)

make_requirements.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+import os
+import ast
+import sys
+from importlib import metadata
+
+# --- CONFIGURE THIS ---
+PROJECT_PATH = r"C:\Users\santi\Desktop\Oto\Vertebra-Landmark-Detection"
+OUTPUT_FILE  = os.path.join(PROJECT_PATH, "requirements.txt")
+# ----------------------
+
+def find_py_files(root):
+    for dirpath, dirnames, filenames in os.walk(root):
+        # skip __pycache__
+        dirnames[:] = [d for d in dirnames if d != "__pycache__"]
+        for fname in filenames:
+            if fname.endswith(".py"):
+                yield os.path.join(dirpath, fname)
+
+def collect_imports(py_path):
+    with open(py_path, "r", encoding="utf8") as f:
+        node = ast.parse(f.read(), filename=py_path)
+    imports = set()
+    for stmt in ast.walk(node):
+        if isinstance(stmt, ast.Import):
+            for n in stmt.names:
+                imports.add(n.name.split(".")[0])
+        elif isinstance(stmt, ast.ImportFrom):
+            if stmt.module and stmt.level == 0:
+                imports.add(stmt.module.split(".")[0])
+    return imports
+
+def is_local_module(mod_name, project_root):
+    # if there's a folder or file matching mod_name in project, treat as local
+    path1 = os.path.join(project_root, mod_name + ".py")
+    path2 = os.path.join(project_root, mod_name)
+    return os.path.exists(path1) or os.path.exists(path2)
+
+def main():
+    all_imports = set()
+    for py in find_py_files(PROJECT_PATH):
+        all_imports |= collect_imports(py)
+
+    # filter out builtins, stdlib, and local modules
+    externals = set()
+    for mod in sorted(all_imports):
+        if is_local_module(mod, PROJECT_PATH):
+            continue
+        try:
+            # try to see if it's installed as a distribution
+            dist = metadata.distribution(mod)
+            externals.add(f"{dist.metadata['Name']}=={dist.version}")
+        except metadata.PackageNotFoundError:
+            # not a top-level distribution, maybe stdlib or nested import
+            # skip modules that come with the stdlib
+            # crude check: if we can import and it's in stdlib path, skip
+            try:
+                m = __import__(mod)
+                if hasattr(m, "__file__") and "site-packages" in (m.__file__ or ""):
+                    # lives in site-packages but dist metadata missing: include without version
+                    externals.add(mod)
+            except ImportError:
+                pass
+
+    # write requirements.txt
+    with open(OUTPUT_FILE, "w", encoding="utf8") as out:
+        for line in sorted(externals):
+            out.write(line + "\n")
+
+    print(f"Written {len(externals)} packages to {OUTPUT_FILE}")
+
+if __name__ == "__main__":
+    main()

pre_proc.py

@@ -0,0 +1,146 @@
+import cv2
+import torch
+from draw_gaussian import *
+import transform
+import math
+
+
+def processing_test(image, input_h, input_w):
+    image = cv2.resize(image, (input_w, input_h))
+    out_image = image.astype(np.float32) / 255.
+    out_image = out_image - 0.5
+    out_image = out_image.transpose(2, 0, 1).reshape(1, 3, input_h, input_w)
+    out_image = torch.from_numpy(out_image)
+    return out_image
+
+
+def draw_spinal(pts, out_image):
+    colors = [(0, 0, 255), (0, 255, 255), (255, 0, 255), (0, 255, 0)]
+    for i in range(4):
+        cv2.circle(out_image, (int(pts[i, 0]), int(pts[i, 1])), 3, colors[i], 1, 1)
+        cv2.putText(out_image, '{}'.format(i+1), (int(pts[i, 0]), int(pts[i, 1])),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0,0,0),1,1)
+    for i,j in zip([0,1,2,3], [1,2,3,0]):
+        cv2.line(out_image,
+                 (int(pts[i, 0]), int(pts[i, 1])),
+                 (int(pts[j, 0]), int(pts[j, 1])),
+                 color=colors[i], thickness=1, lineType=1)
+    return out_image
+
+
+def rearrange_pts(pts):
+    # rearrange left right sequence
+    boxes = []
+    centers = []
+    for k in range(0, len(pts), 4):
+        pts_4 = pts[k:k+4,:]
+        x_inds = np.argsort(pts_4[:, 0])
+        pt_l = np.asarray(pts_4[x_inds[:2], :])
+        pt_r = np.asarray(pts_4[x_inds[2:], :])
+        y_inds_l = np.argsort(pt_l[:,1])
+        y_inds_r = np.argsort(pt_r[:,1])
+        tl = pt_l[y_inds_l[0], :]
+        bl = pt_l[y_inds_l[1], :]
+        tr = pt_r[y_inds_r[0], :]
+        br = pt_r[y_inds_r[1], :]
+        # boxes.append([tl, tr, bl, br])
+        boxes.append(tl)
+        boxes.append(tr)
+        boxes.append(bl)
+        boxes.append(br)
+        centers.append(np.mean(pts_4, axis=0))
+    bboxes = np.asarray(boxes, np.float32)
+    # rearrange top to bottom sequence
+    centers = np.asarray(centers, np.float32)
+    sort_tb = np.argsort(centers[:,1])
+    new_bboxes = []
+    for sort_i in sort_tb:
+        new_bboxes.append(bboxes[4*sort_i, :])
+        new_bboxes.append(bboxes[4*sort_i+1, :])
+        new_bboxes.append(bboxes[4*sort_i+2, :])
+        new_bboxes.append(bboxes[4*sort_i+3, :])
+    new_bboxes = np.asarray(new_bboxes, np.float32)
+    return new_bboxes
+
+
+def generate_ground_truth(image,
+                          pts_2,
+                          image_h,
+                          image_w,
+                          img_id):
+    hm = np.zeros((1, image_h, image_w), dtype=np.float32)
+    wh = np.zeros((17, 2*4), dtype=np.float32)
+    reg = np.zeros((17, 2), dtype=np.float32)
+    ind = np.zeros((17), dtype=np.int64)
+    reg_mask = np.zeros((17), dtype=np.uint8)
+
+    if pts_2[:,0].max()>image_w:
+        print('w is big', pts_2[:,0].max())
+    if pts_2[:,1].max()>image_h:
+        print('h is big', pts_2[:,1].max())
+
+    if pts_2.shape[0]!=68:
+        print('ATTENTION!! image {} pts does not equal to 68!!! '.format(img_id))
+
+    for k in range(17):
+        pts = pts_2[4*k:4*k+4,:]
+        bbox_h = np.mean([np.sqrt(np.sum((pts[0,:]-pts[2,:])**2)),
+                          np.sqrt(np.sum((pts[1,:]-pts[3,:])**2))])
+        bbox_w = np.mean([np.sqrt(np.sum((pts[0,:]-pts[1,:])**2)),
+                          np.sqrt(np.sum((pts[2,:]-pts[3,:])**2))])
+        cen_x, cen_y = np.mean(pts, axis=0)
+        ct = np.asarray([cen_x, cen_y], dtype=np.float32)
+        ct_int = ct.astype(np.int32)
+        radius = gaussian_radius((math.ceil(bbox_h), math.ceil(bbox_w)))
+        radius = max(0, int(radius))
+        draw_umich_gaussian(hm[0,:,:], ct_int, radius=radius)
+        ind[k] = ct_int[1] * image_w + ct_int[0]
+        reg[k] = ct - ct_int
+        reg_mask[k] = 1
+        for i in range(4):
+            wh[k,2*i:2*i+2] = ct-pts[i,:]
+
+    ret = {'input': image,
+           'hm': hm,
+           'ind': ind,
+           'reg': reg,
+           'wh': wh,
+           'reg_mask': reg_mask,
+           }
+
+    return ret
+
+# def filter_pts(pts, w, h):
+#     pts_new = []
+#     for pt in pts:
+#         if any(pt) < 0 or pt[0] > w - 1 or pt[1] > h - 1:
+#             continue
+#         else:
+#             pts_new.append(pt)
+#     return np.asarray(pts_new, np.float32)
+
+
+def processing_train(image, pts, image_h, image_w, down_ratio, aug_label, img_id):
+    # filter pts ----------------------------------------------------
+    h,w,c = image.shape
+    # pts = filter_pts(pts, w, h)
+    # ---------------------------------------------------------------
+    data_aug = {'train': transform.Compose([transform.ConvertImgFloat(),
+                                            transform.PhotometricDistort(),
+                                            transform.Expand(max_scale=1.5, mean=(0, 0, 0)),
+                                            transform.RandomMirror_w(),
+                                            transform.Resize(h=image_h, w=image_w)]),
+                'val': transform.Compose([transform.ConvertImgFloat(),
+                                          transform.Resize(h=image_h, w=image_w)])}
+    if aug_label:
+        out_image, pts = data_aug['train'](image.copy(), pts)
+    else:
+        out_image, pts = data_aug['val'](image.copy(), pts)
+
+    out_image = np.clip(out_image, a_min=0., a_max=255.)
+    out_image = np.transpose(out_image / 255. - 0.5, (2,0,1))
+    pts = rearrange_pts(pts)
+    pts2 = transform.rescale_pts(pts, down_ratio=down_ratio)
+
+    return np.asarray(out_image, np.float32), pts2
+

test.py

@@ -0,0 +1,123 @@
+import torch
+import numpy as np
+from models import spinal_net
+import cv2
+import decoder
+import os
+from dataset import BaseDataset
+import draw_points
+
+def apply_mask(image, mask, alpha=0.5):
+    """Apply the given mask to the image.
+    """
+    color = np.random.rand(3)
+    for c in range(3):
+        image[:, :, c] = np.where(mask == 1,
+                                  image[:, :, c] *
+                                  (1 - alpha) + alpha * color[c] * 255,
+                                  image[:, :, c])
+    return image
+
+class Network(object):
+    def __init__(self, args):
+        torch.manual_seed(317)
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        heads = {'hm': args.num_classes,
+                 'reg': 2*args.num_classes,
+                 'wh': 2*4,}
+
+        self.model = spinal_net.SpineNet(heads=heads,
+                                         pretrained=True,
+                                         down_ratio=args.down_ratio,
+                                         final_kernel=1,
+                                         head_conv=256)
+        self.num_classes = args.num_classes
+        self.decoder = decoder.DecDecoder(K=args.K, conf_thresh=args.conf_thresh)
+        self.dataset = {'spinal': BaseDataset}
+
+    def load_model(self, model, resume):
+        checkpoint = torch.load(resume, map_location=lambda storage, loc: storage)
+        print('loaded weights from {}, epoch {}'.format(resume, checkpoint['epoch']))
+        state_dict_ = checkpoint['state_dict']
+        model.load_state_dict(state_dict_, strict=False)
+        return model
+
+    def map_mask_to_image(self, mask, img, color=None):
+        if color is None:
+            color = np.random.rand(3)
+        mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
+        mskd = img * mask
+        clmsk = np.ones(mask.shape) * mask
+        clmsk[:, :, 0] = clmsk[:, :, 0] * color[0] * 256
+        clmsk[:, :, 1] = clmsk[:, :, 1] * color[1] * 256
+        clmsk[:, :, 2] = clmsk[:, :, 2] * color[2] * 256
+        img = img + 1. * clmsk - 1. * mskd
+        return np.uint8(img)
+
+
+    def test(self, args, save):
+        save_path = 'weights_'+args.dataset
+        self.model = self.load_model(self.model, os.path.join(save_path, args.resume))
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        dataset_module = self.dataset[args.dataset]
+        dsets = dataset_module(data_dir=args.data_dir,
+                               phase='test',
+                               input_h=args.input_h,
+                               input_w=args.input_w,
+                               down_ratio=args.down_ratio)
+
+        data_loader = torch.utils.data.DataLoader(dsets,
+                                                  batch_size=1,
+                                                  shuffle=False,
+                                                  num_workers=1,
+                                                  pin_memory=True)
+
+
+        for cnt, data_dict in enumerate(data_loader):
+            images = data_dict['images'][0]
+            img_id = data_dict['img_id'][0]
+            images = images.to('cuda')
+            print('processing {}/{} image ... {}'.format(cnt, len(data_loader), img_id))
+            with torch.no_grad():
+                output = self.model(images)
+                hm = output['hm']
+                wh = output['wh']
+                reg = output['reg']
+
+            torch.cuda.synchronize(self.device)
+            pts2 = self.decoder.ctdet_decode(hm, wh, reg)   # 17, 11
+            pts0 = pts2.copy()
+            pts0[:,:10] *= args.down_ratio
+
+            print('totol pts num is {}'.format(len(pts2)))
+
+            ori_image = dsets.load_image(dsets.img_ids.index(img_id))
+            ori_image_regress = cv2.resize(ori_image, (args.input_w, args.input_h))
+            ori_image_points = ori_image_regress.copy()
+
+            h,w,c = ori_image.shape
+            pts0 = np.asarray(pts0, np.float32)
+            # pts0[:,0::2] = pts0[:,0::2]/args.input_w*w
+            # pts0[:,1::2] = pts0[:,1::2]/args.input_h*h
+            sort_ind = np.argsort(pts0[:,1])
+            pts0 = pts0[sort_ind]
+
+            ori_image_regress, ori_image_points = draw_points.draw_landmarks_regress_test(pts0,
+                                                                                          ori_image_regress,
+                                                                                          ori_image_points)
+
+            if save:
+                # 1) กำหนดโฟลเดอร์ผลลัพธ์
+                save_dir = os.path.join('results_'+args.dataset)
+                os.makedirs(save_dir, exist_ok=True)
+
+                # 2) บันทึกพิกัดลง .txt
+                txt_path = os.path.join(save_dir, f'{img_id}.txt')
+                # สมมติ pts0 เป็น array shape (N,2) หรือ (N,4) ตามที่คุณอยากบันทึก
+                np.savetxt(txt_path, pts0, fmt='%.4f')
+
+                # 3) บันทึกภาพ overlay
+                img_path = os.path.join(save_dir, f'{img_id}_pred.jpg')
+                cv2.imwrite(img_path, ori_image_points)

train.py

@@ -0,0 +1,169 @@
+import torch
+import torch.nn as nn
+import os
+import numpy as np
+from models import spinal_net
+import decoder
+import loss
+from dataset import BaseDataset
+
+def collater(data):
+    out_data_dict = {}
+    for name in data[0]:
+        out_data_dict[name] = []
+    for sample in data:
+        for name in sample:
+            out_data_dict[name].append(torch.from_numpy(sample[name]))
+    for name in out_data_dict:
+        out_data_dict[name] = torch.stack(out_data_dict[name], dim=0)
+    return out_data_dict
+
+class Network(object):
+    def __init__(self, args):
+        torch.manual_seed(317)
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        heads = {'hm': args.num_classes,
+                 'reg': 2*args.num_classes,
+                 'wh': 2*4,}
+
+        self.model = spinal_net.SpineNet(heads=heads,
+                                         pretrained=True,
+                                         down_ratio=args.down_ratio,
+                                         final_kernel=1,
+                                         head_conv=256)
+        self.num_classes = args.num_classes
+        self.decoder = decoder.DecDecoder(K=args.K, conf_thresh=args.conf_thresh)
+        self.dataset = {'spinal': BaseDataset}
+
+
+    def save_model(self, path, epoch, model):
+        if isinstance(model, torch.nn.DataParallel):
+            state_dict = model.module.state_dict()
+        else:
+            state_dict = model.state_dict()
+        data = {'epoch': epoch, 'state_dict': state_dict}
+        torch.save(data, path)
+
+    def load_model(self, model, resume, strict=True):
+        checkpoint = torch.load(resume, map_location=lambda storage, loc: storage)
+        print('loaded weights from {}, epoch {}'.format(resume, checkpoint['epoch']))
+        state_dict_ = checkpoint['state_dict']
+        state_dict = {}
+
+        for k in state_dict_:
+            if k.startswith('module') and not k.startswith('module_list'):
+                state_dict[k[7:]] = state_dict_[k]
+            else:
+                state_dict[k] = state_dict_[k]
+        model_state_dict = model.state_dict()
+
+        if not strict:
+            for k in state_dict:
+                if k in model_state_dict:
+                    if state_dict[k].shape != model_state_dict[k].shape:
+                        print('Skip loading parameter {}, required shape{}, ' \
+                              'loaded shape{}.'.format(k, model_state_dict[k].shape, state_dict[k].shape))
+                        state_dict[k] = model_state_dict[k]
+                else:
+                    print('Drop parameter {}.'.format(k))
+            for k in model_state_dict:
+                if not (k in state_dict):
+                    print('No param {}.'.format(k))
+                    state_dict[k] = model_state_dict[k]
+        model.load_state_dict(state_dict, strict=False)
+        return model
+
+    def train_network(self, args):
+        save_path = 'weights_'+args.dataset
+        if not os.path.exists(save_path):
+            os.mkdir(save_path)
+        self.optimizer = torch.optim.Adam(self.model.parameters(), args.init_lr)
+        scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.96, last_epoch=-1)
+        if args.ngpus>0:
+            if torch.cuda.device_count() > 1:
+                print("Let's use", torch.cuda.device_count(), "GPUs!")
+                self.model = nn.DataParallel(self.model)
+
+        self.model.to(self.device)
+
+        criterion = loss.LossAll()
+        print('Setting up data...')
+
+        dataset_module = self.dataset[args.dataset]
+
+        dsets = {x: dataset_module(data_dir=args.data_dir,
+                                   phase=x,
+                                   input_h=args.input_h,
+                                   input_w=args.input_w,
+                                   down_ratio=args.down_ratio)
+                 for x in ['train', 'val']}
+
+        dsets_loader = {'train': torch.utils.data.DataLoader(dsets['train'],
+                                                             batch_size=args.batch_size,
+                                                             shuffle=True,
+                                                             num_workers=args.num_workers,
+                                                             pin_memory=True,
+                                                             drop_last=True,
+                                                             collate_fn=collater),
+
+                        'val':torch.utils.data.DataLoader(dsets['val'],
+                                                          batch_size=1,
+                                                          shuffle=False,
+                                                          num_workers=1,
+                                                          pin_memory=True,
+                                                          collate_fn=collater)}
+
+
+        print('Starting training...')
+        train_loss = []
+        val_loss = []
+        for epoch in range(1, args.num_epoch+1):
+            print('-'*10)
+            print('Epoch: {}/{} '.format(epoch, args.num_epoch))
+            epoch_loss = self.run_epoch(phase='train',
+                                        data_loader=dsets_loader['train'],
+                                        criterion=criterion)
+            train_loss.append(epoch_loss)
+            scheduler.step(epoch)
+
+            epoch_loss = self.run_epoch(phase='val',
+                                        data_loader=dsets_loader['val'],
+                                        criterion=criterion)
+            val_loss.append(epoch_loss)
+
+            np.savetxt(os.path.join(save_path, 'train_loss.txt'), train_loss, fmt='%.6f')
+            np.savetxt(os.path.join(save_path, 'val_loss.txt'), val_loss, fmt='%.6f')
+
+            if epoch % 10 == 0 or epoch ==1:
+                self.save_model(os.path.join(save_path, 'model_{}.pth'.format(epoch)), epoch, self.model)
+
+            if len(val_loss)>1:
+                if val_loss[-1]<np.min(val_loss[:-1]):
+                    self.save_model(os.path.join(save_path, 'model_last.pth'), epoch, self.model)
+
+    def run_epoch(self, phase, data_loader, criterion):
+        if phase == 'train':
+            self.model.train()
+        else:
+            self.model.eval()
+        running_loss = 0.
+        for data_dict in data_loader:
+            for name in data_dict:
+                data_dict[name] = data_dict[name].to(device=self.device)
+            if phase == 'train':
+                self.optimizer.zero_grad()
+                with torch.enable_grad():
+                    pr_decs = self.model(data_dict['input'])
+                    loss = criterion(pr_decs, data_dict)
+                    loss.backward()
+                    self.optimizer.step()
+            else:
+                with torch.no_grad():
+                    pr_decs = self.model(data_dict['input'])
+                    loss = criterion(pr_decs, data_dict)
+
+            running_loss += loss.item()
+        epoch_loss = running_loss / len(data_loader)
+        print('{} loss: {}'.format(phase, epoch_loss))
+        return epoch_loss
+

transform.py

@@ -0,0 +1,181 @@
+
+import numpy as np
+from numpy import random
+import cv2
+
+
+def rescale_pts(pts, down_ratio):
+    return np.asarray(pts, np.float32)/float(down_ratio)
+
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, img, pts):
+        for t in self.transforms:
+            img, pts = t(img, pts)
+        return img, pts
+
+class ConvertImgFloat(object):
+    def __call__(self, img, pts):
+        return img.astype(np.float32), pts.astype(np.float32)
+
+class RandomContrast(object):
+    def __init__(self, lower=0.5, upper=1.5):
+        self.lower = lower
+        self.upper = upper
+        assert self.upper >= self.lower, "contrast upper must be >= lower."
+        assert self.lower >= 0, "contrast lower must be non-negative."
+
+    def __call__(self, img, pts):
+        if random.randint(2):
+            alpha = random.uniform(self.lower, self.upper)
+            img *= alpha
+        return img, pts
+
+
+class RandomBrightness(object):
+    def __init__(self, delta=32):
+        assert delta >= 0.0
+        assert delta <= 255.0
+        self.delta = delta
+
+    def __call__(self, img, pts):
+        if random.randint(2):
+            delta = random.uniform(-self.delta, self.delta)
+            img += delta
+        return img, pts
+
+class SwapChannels(object):
+    def __init__(self, swaps):
+        self.swaps = swaps
+    def __call__(self, img):
+        img = img[:, :, self.swaps]
+        return img
+
+
+class RandomLightingNoise(object):
+    def __init__(self):
+        self.perms = ((0, 1, 2), (0, 2, 1),
+                      (1, 0, 2), (1, 2, 0),
+                      (2, 0, 1), (2, 1, 0))
+    def __call__(self, img, pts):
+        if random.randint(2):
+            swap = self.perms[random.randint(len(self.perms))]
+            shuffle = SwapChannels(swap)
+            img = shuffle(img)
+        return img, pts
+
+
+class PhotometricDistort(object):
+    def __init__(self):
+        self.pd = RandomContrast()
+        self.rb = RandomBrightness()
+        self.rln = RandomLightingNoise()
+
+    def __call__(self, img, pts):
+        img, pts = self.rb(img, pts)
+        if random.randint(2):
+            distort = self.pd
+        else:
+            distort = self.pd
+        img, pts = distort(img, pts)
+        img, pts = self.rln(img, pts)
+        return img, pts
+
+
+class Expand(object):
+    def __init__(self, max_scale = 1.5, mean = (0.5, 0.5, 0.5)):
+        self.mean = mean
+        self.max_scale = max_scale
+
+    def __call__(self, img, pts):
+        if random.randint(2):
+            return img, pts
+        h,w,c = img.shape
+        ratio = random.uniform(1,self.max_scale)
+        y1 = random.uniform(0, h*ratio-h)
+        x1 = random.uniform(0, w*ratio-w)
+        if np.max(pts[:,0])+int(x1)>w-1 or np.max(pts[:,1])+int(y1)>h-1:  # keep all the pts
+            return img, pts
+        else:
+            expand_img = np.zeros(shape=(int(h*ratio), int(w*ratio),c),dtype=img.dtype)
+            expand_img[:,:,:] = self.mean
+            expand_img[int(y1):int(y1+h), int(x1):int(x1+w)] = img
+            pts[:, 0] += int(x1)
+            pts[:, 1] += int(y1)
+            return expand_img, pts
+
+
+class RandomSampleCrop(object):
+    def __init__(self, ratio=(0.5, 1.5), min_win = 0.9):
+        self.sample_options = (
+            # using entire original input image
+            None,
+            # sample a patch s.t. MIN jaccard w/ obj in .1,.3,.4,.7,.9
+            # (0.1, None),
+            # (0.3, None),
+            (0.7, None),
+            (0.9, None),
+            # randomly sample a patch
+            (None, None),
+        )
+        self.ratio = ratio
+        self.min_win = min_win
+
+    def __call__(self, img, pts):
+        height, width ,_ = img.shape
+        while True:
+            mode = random.choice(self.sample_options)
+            if mode is None:
+                return img, pts
+            for _ in range(50):
+                current_img = img
+                current_pts = pts
+                w = random.uniform(self.min_win*width, width)
+                h = random.uniform(self.min_win*height, height)
+                if h/w<self.ratio[0] or h/w>self.ratio[1]:
+                    continue
+                y1 = random.uniform(height-h)
+                x1 = random.uniform(width-w)
+                rect = np.array([int(y1), int(x1), int(y1+h), int(x1+w)])
+                current_img = current_img[rect[0]:rect[2], rect[1]:rect[3], :]
+                current_pts[:, 0] -= rect[1]
+                current_pts[:, 1] -= rect[0]
+                pts_new = []
+                for pt in current_pts:
+                    if any(pt)<0 or pt[0]>current_img.shape[1]-1 or pt[1]>current_img.shape[0]-1:
+                        continue
+                    else:
+                        pts_new.append(pt)
+
+                return current_img, np.asarray(pts_new, np.float32)
+
+class RandomMirror_w(object):
+    def __call__(self, img, pts):
+        _,w,_ = img.shape
+        if random.randint(2):
+            img = img[:,::-1,:]
+            pts[:,0] = w-pts[:,0]
+        return img, pts
+
+class RandomMirror_h(object):
+    def __call__(self, img, pts):
+        h,_,_ = img.shape
+        if random.randint(2):
+            img = img[::-1,:,:]
+            pts[:,1] = h-pts[:,1]
+        return img, pts
+
+
+class Resize(object):
+    def __init__(self, h, w):
+        self.dsize = (w,h)
+
+    def __call__(self, img, pts):
+        h,w,c = img.shape
+        pts[:, 0] = pts[:, 0]/w*self.dsize[0]
+        pts[:, 1] = pts[:, 1]/h*self.dsize[1]
+        img = cv2.resize(img, dsize=self.dsize)
+        return img, np.asarray(pts)