Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,86 @@
-import altair as alt
+import streamlit as st 
 import numpy as np
-import pandas as pd
-import streamlit as st
-
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+from PIL import Image
+import cv2
+from ultralytics import YOLO
+
+@st.cache_resource(show_spinner=False)
+def load_model():
+    return YOLO('best.pt')  # Make sure the model file is in the same folder or provide full path
+
+def predict(image, model):
+    img_np = np.array(image)
+    if img_np.shape[2] == 4:
+        img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2RGB)
+    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+    results = model(img_bgr)
+    return results[0]
+
+def draw_results(image, results):
+    img_annotated = results.plot()
+    img_annotated_rgb = cv2.cvtColor(img_annotated, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(img_annotated_rgb)
+
+def crop_box(image, box):
+    img_np = np.array(image)
+    x1, y1, x2, y2 = map(int, box.xyxy[0])
+    h, w = img_np.shape[:2]
+    pad_x, pad_y = int((x2 - x1) * 0.1), int((y2 - y1) * 0.1)
+    x1, y1 = max(0, x1 - pad_x), max(0, y1 - pad_y)
+    x2, y2 = min(w, x2 + pad_x), min(h, y2 + pad_y)
+    cropped = img_np[y1:y2, x1:x2]
+    return Image.fromarray(cropped)
+
+def main():
+    st.title("🦷 Interactive Teeth Segmentation & Annotation")
+
+    st.markdown("""
+    Upload a panoramic dental X-ray image.  
+    After detection, you can edit the tooth number labels for each detected tooth.
+    """)
+
+    model = load_model()
+    uploaded_file = st.file_uploader("Upload panoramic dental X-ray", type=["png", "jpg", "jpeg"])
+
+    if uploaded_file:
+        image = Image.open(uploaded_file).convert("RGB")
+        with st.spinner("Detecting teeth..."):
+            results = predict(image, model)
+
+        annotated_img = draw_results(image, results)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.subheader("Original Image")
+            st.image(image, use_column_width=True)
+        with col2:
+            st.subheader("Detection & Segmentation")
+            st.image(annotated_img, use_column_width=True)
+
+        if len(results.boxes) > 0:
+            st.markdown("### Edit Detected Tooth Numbers")
+
+            edited_labels = {}
+
+            for i, box in enumerate(results.boxes):
+                class_id = int(box.cls[0])
+                default_label = results.names[class_id]
+                confidence = box.conf[0].item()
+                cropped = crop_box(image, box)
+
+                st.markdown(f"**Tooth {i+1}** (Confidence: {confidence:.2%})")
+                st.image(cropped, width=150)
+
+                new_label = st.text_input(f"Change tooth number (default: {default_label})",
+                                          value=default_label,
+                                          key=f"label_{i}")
+
+                edited_labels[i] = new_label
+
+            st.markdown("---")
+            st.subheader("Final Tooth Labels")
+            for i, label in edited_labels.items():
+                st.write(f"Tooth {i+1}: {label}")
+
+        else:
+            st.info("No teeth detected in this image.")