Update app.py

app.py

@@ -88,8 +88,8 @@ 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"
+        ["How to use", "AP - Detection", "LA - Image Segmetation", "Contract"],
+        index=3,  # default to "AP"
         help="Choose which view to display"
     )
 
@@ -98,14 +98,25 @@ if feature == "How to use":
 
     col1, col2, col3 = st.columns(3)
 
+    card_style = """
+        border:2px solid #00BFFF;
+        border-radius:10px;
+        padding:15px;
+        text-align:center;
+        background-color:#F0F8FF;
+    """
+
+    title_style = "color:#000f14; margin-bottom:10px;"
+    body_style  = "color:#000f14; text-align:left;"
+
     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>
+            f"""
+            <div style="{card_style}">
+                <h2 style="{title_style}">Step 1️⃣</h2>
+                <p style="{body_style}">Go to <b>AP - Detection</b> or <b>LA - Image Segmentation</b></p>
+                <p style="{body_style}">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
@@ -113,12 +124,12 @@ if feature == "How to use":
 
     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>
+            f"""
+            <div style="{card_style}">
+                <h2 style="{title_style}">Step 2️⃣</h2>
+                <p style="{body_style}">Press the <b>Enter</b> button.</p>
+                <p style="{body_style}">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
@@ -126,12 +137,12 @@ if feature == "How to use":
 
     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>
+            f"""
+            <div style="{card_style}">
+                <h2 style="{title_style}">Step 3️⃣</h2>
+                <p style="{body_style}">See the prediction results:</p>
+                <p style="{body_style}">1. Bounding boxes & landmarks (AP)</p>
+                <p style="{body_style}">2. Segmentation masks (LA)</p>
             </div>
             """,
             unsafe_allow_html=True
@@ -140,18 +151,17 @@ if feature == "How to use":
     st.markdown(" ")
     st.info("สามารถเลือกฟีเจอร์ได้ผ่าน Select Feature โดยแต่ล่ะฟีเจอร์จะมีตัวอย่างกำกับให้ว่าเป็นยังไง")
 
-# store original dimensions
+# … (any code above)
+
 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"
@@ -162,40 +172,34 @@ elif feature == "AP - Detection":
         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
-
         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)
-
+                     use_container_width=True)
         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!")
+                st.error(f"Cannot find {sample_img}")
 
     with col5:
         st.subheader("2️⃣ Predictions")
     with col6:
         st.subheader("3️⃣ Heatmap")
 
-    # ─── ARGS & CHECKPOINT ─────────────────────────────────
     args = Namespace(
-        resume="model_30.pth",
+        resume="model_50.pth",
         data_dir=os.path.join(REPO, "dataPath"),
         dataset="spinal",
         phase="test",
@@ -213,21 +217,16 @@ elif feature == "AP - Detection":
     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"
+        name = (os.path.splitext(uploaded.name)[0]
+                if uploaded else os.path.splitext(sample_img)[0]) + ".jpg"
+        test_dir = os.path.join(args.data_dir, "data", "test")
+        os.makedirs(test_dir, exist_ok=True)
+        cv2.imwrite(os.path.join(test_dir, name), img0)
 
-        testd = os.path.join(args.data_dir, "data", "test")
-        os.makedirs(testd, exist_ok=True)
-        cv2.imwrite(os.path.join(testd, name), img0)
-
-        # patch BaseDataset to only load our one image
         orig_init = BaseDataset.__init__
-        def patched_init(self, data_dir, phase, input_h=None, input_w=None, down_ratio=4):
+        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]
@@ -238,84 +237,117 @@ elif feature == "AP - Detection":
             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]
+        pred_file = next(
+            f for f in os.listdir(out_dir)
+            if f.startswith(name) and f.endswith("_pred.jpg")
+        )
         txtf = os.path.join(out_dir, f"{name}.txt")
         imgf = os.path.join(out_dir, pred_file)
 
-        # ─── Annotated Predictions ─────────────────────────
-        base = cv2.imread(imgf)
+        # ─── Annotated predictions ─────────────────────────────────────
+        ann = 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)
-
-        # compute mid‐points and heights
-        cts, heights = [], []
-        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])
-            cts.append((int(tm[0]), int((tm[1]+bm[1])//2)))
-            heights.append(int(bm[1] - tm[1]))
-
-        # draw lines on 'ann'
-        ann = base.copy()
-        for (cx, cy), h in zip(cts, heights):
-            # top‐bottom line
-            cv2.line(ann, (cx, cy - h//2), (cx, cy + h//2), (0, 255, 255), 2)
-
-        # neighbor‐based compression percentages
-        for idx, ((cx, cy), height) in enumerate(zip(cts, heights)):
-            # reference = average of neighbors
-            if 0 < idx < len(heights) - 1:
-                ref_h = (heights[idx - 1] + heights[idx + 1]) / 2
-            else:
-                ref_h = np.median(heights)
-
-            percent = abs((ref_h - height) / ref_h * 100)
-
-            # color thresholds
-            if percent > 40:
-                color = (0, 0, 255)
-            elif percent > 20:
-                color = (0, 165, 255)
-            else:
-                color = (0, 255, 0)
-
-            # label
-            text_pos = (cx + 5, cy)
-            cv2.putText(ann, f"{percent:.0f}%", text_pos,
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2, cv2.LINE_AA)
-            print(f"กระดูกตัวที่ {idx+1}: Compression = {percent:.1f}%")
+        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)
+
+        for x1, y1, x2, y2 in zip(tlx, tly, trx, try_):
+            cv2.line(ann, (x1, y1), (x2, y2), (255,255,0), 2)
+
+        for x1,y1,x2,y2,x3,y3,x4,y4 in zip(
+            tlx, tly, trx, try_, blx, bly, brx, bry
+        ):
+            top_mid = np.array([(x1+x2)/2, (y1+y2)/2])
+            bot_mid = np.array([(x3+x4)/2, (y3+y4)/2])
+            p0 = tuple(top_mid.astype(int))
+            p1 = tuple(bot_mid.astype(int))
+            cv2.line(ann, p0, p1, (0,255,255), 2)
+
+            h_before = np.linalg.norm(bot_mid - top_mid)
+            h_after  = 2 * int(h_before * 0.4)
+            pct      = ((h_before - h_after) / h_before * 100) - 10
+            clr = (0,0,255) if pct > 40 else (
+                  (0,165,255) if pct > 20 else (0,255,255))
+            text_pos = (x2 + 5, y2 - 5)
+            cv2.putText(
+                ann, f"{pct:.0f}%", text_pos,
+                cv2.FONT_HERSHEY_SIMPLEX, 0.5, clr, 2, cv2.LINE_AA
+            )
 
-        # show annotated image
-        ann_resized = cv2.resize(ann, (orig_w, orig_h),
-                                 interpolation=cv2.INTER_LINEAR)
+        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)
+            st.image(
+                cv2.cvtColor(ann_resized, cv2.COLOR_BGR2RGB),
+                use_container_width=True
+            )
 
-        # ─── Heatmap ──────────────────────────────────────────
+        # ─── Heatmap overlay + connecting lines ─────────────────────────
+        base = cv2.imread(imgf)
         H, W = base.shape[:2]
         heat = np.zeros((H, W), np.float32)
+        cts = []
+        for (x1, y1), (x2, y2) in zip(zip(tlx, tly), zip(trx, try_)):
+            tm = np.array([(x1 + x2)/2, (y1 + y2)/2])
+            cts.append((int(tm[0]), int(tm[1])))
+
         for cx, cy in cts:
             blob = np.zeros_like(heat)
             blob[cy, cx] = 1.0
-            heat += cv2.GaussianBlur(blob, (0, 0), sigmaX=8, sigmaY=8)
+            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)
+
+        for p1, p2 in zip(cts, cts[1:]):
+            cv2.line(overlay, p1, p2, (0,255,255), 2)
+
+        # ─── Cobb‑angle original logic ────────────────────────────────
+        vecs = np.diff(np.array(cts), axis=0)
+        angles = np.degrees(np.arctan2(vecs[:,1], vecs[:,0]))
+        idx_max = int(np.argmax(angles))
+        idx_min = int(np.argmin(angles))
+        cobb = abs(angles[idx_max] - angles[idx_min])
+
+        # ─── highlight apex of curvature ─────────────────────────────
+        # compute local curvature angles
+        norms = np.linalg.norm(vecs, axis=1, keepdims=True)
+        unit = vecs / norms
+        dots = np.sum(unit[:-1] * unit[1:], axis=1)
+        dots = np.clip(dots, -1.0, 1.0)
+        thetas = np.degrees(np.arccos(dots))
+        apex_idx = int(np.argmax(thetas)) + 1  # vertex index
+        vx, vy = cts[apex_idx]
+        cv2.circle(overlay, (vx, vy), 15, (0, 0, 255), 2)
+
+        # ─── draw centered Cobb text ────────────────────────────────
+        text1 = "Cobb Angle"
+        text2 = f"{cobb:.1f}"
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        scale, thickness = 1.0, 2
+        (w1,h1),_ = cv2.getTextSize(text1, font, scale, thickness)
+        (w2,h2),_ = cv2.getTextSize(text2, font, scale, thickness)
+        x1 = (W - w1)//2; y1 = H//2 - h1 - 10
+        x2 = (W - w2)//2; y2 = H//2 + h2 + 10
+        cv2.putText(overlay, text1, (x1, y1), font, scale, (0,255,255), thickness, cv2.LINE_AA)
+        cv2.putText(overlay, text2, (x2, y2), font, scale, (0,255,255), thickness, cv2.LINE_AA)
+
+        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)
+            st.image(
+                cv2.cvtColor(overlay_resized, cv2.COLOR_BGR2RGB),
+                use_container_width=True
+            )
+
 
 
 elif feature == "LA - Image Segmetation":
@@ -368,7 +400,7 @@ elif feature == "LA - Image Segmetation":
         # ─── PREDICTION ────────────────────────────────────
         if img0 is not None and run_la:
             img_np = np.array(img0)
-            model = YOLO('./best.pt')  # path to your weights
+            model = YOLO('./best_100.pt')  # path to your weights
             with st.spinner("Running YOLO model…"):
                 results = model(img_np, imgsz=640)
 
@@ -392,40 +424,59 @@ elif feature == "LA - Image Segmetation":
 
 
 elif feature == "Contract":
+    # shared styles
+    card_style = """
+        border:2px solid #0080FF;
+        border-radius:10px;
+        padding:15px;
+        text-align:center;
+        background-color:#F0F8FF;
+    """
+    title_style = "color:#00BFFF; margin-bottom:8px;"  # names
+    body_style  = "color:#87CEEB; text-decoration:none;"
+
     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
+            f"""
+            <div style="{card_style}">
+                <h3 style="{title_style}">Thitsanapat S.</h3>
+                <a href="https://www.facebook.com/thitsanapat.uma"
+                   target="_blank"
+                   style="{body_style}">
+                   🔗 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
+            f"""
+            <div style="{card_style}">
+                <h3 style="{title_style}">Santipab T.</h3>
+                <a href="https://www.facebook.com/santipab.tongchan.2025"
+                   target="_blank"
+                   style="{body_style}">
+                   🔗 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
+            f"""
+            <div style="{card_style}">
+                <h3 style="{title_style}">Suphanat K.</h3>
+                <a href="https://www.facebook.com/suphanat.kamphapan"
+                   target="_blank"
+                   style="{body_style}">
+                   🔗 Facebook Profile
                 </a>
             </div>
             """,
@@ -433,4 +484,3 @@ elif feature == "Contract":
         )
 
 
-