Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,23 +3,16 @@ import torch
 import torch.nn as nn
 import timm
 import numpy as np
-import cv2
 from PIL import Image
-import io
+import requests
+from io import BytesIO
+import torchvision.transforms as T
+import matplotlib.pyplot as plt
+from huggingface_hub import hf_hub_download
 
-# Hide Streamlit warnings and UI elements
-st.set_page_config(layout="wide")
-st.markdown("""
-    <style>
-        footer {visibility: hidden;}
-    </style>
-""", unsafe_allow_html=True)
-
-MODEL_PATH = "mobilevit_teeth_segmentation.pth"
-
-# === Model Definition ===
+# ========== Model Definition ==========
 class MobileViTSegmentation(nn.Module):
-    def __init__(self, encoder_name='mobilevit_s', pretrained=False):
+    def __init__(self, encoder_name='mobilevit_s', pretrained=True):
         super().__init__()
         self.backbone = timm.create_model(encoder_name, features_only=True, pretrained=pretrained)
         self.encoder_channels = self.backbone.feature_info.channels()
@@ -41,55 +34,49 @@ class MobileViTSegmentation(nn.Module):
         out = nn.functional.interpolate(out, size=(x.shape[2], x.shape[3]), mode='bilinear', align_corners=False)
         return out
 
-# === Load Model ===
+# ========== Load Model ==========
 @st.cache_resource
 def load_model():
-    model = MobileViTSegmentation()
-    state_dict = torch.load( MODEL_PATH, map_location="cpu")
-    model.load_state_dict(state_dict)
+    checkpoint_path = hf_hub_download(repo_id="svsaurav95/ToothSegmentation", filename="mobilevit_teeth_segmentation.pth")
+    model = MobileViTSegmentation(pretrained=False)
+    model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
     model.eval()
     return model
 
 model = load_model()
 
-# === Preprocessing ===
-def preprocess_image(image: Image.Image):
-    image = image.convert("RGB").resize((256, 256))
-    arr = np.array(image).astype(np.float32) / 255.0
-    arr = np.transpose(arr, (2, 0, 1))  # HWC → CHW
-    tensor = torch.tensor(arr).unsqueeze(0)  # Add batch dim
-    return tensor
-
-# === Postprocessing: Overlay Mask ===
-def overlay_mask(image_pil, mask_tensor, threshold=0.7):
-    image = np.array(image_pil.resize((256, 256)))
-    mask = mask_tensor.squeeze().detach().numpy()
-    mask_bin = (mask > threshold).astype(np.uint8) * 255
+# ========== Image Transformation ==========
+transform = T.Compose([
+    T.Resize((256, 256)),
+    T.ToTensor()
+])
 
-    mask_color = np.zeros_like(image)
-    mask_color[..., 2] = mask_bin  # Blue mask
+# ========== Streamlit UI ==========
+st.title("Tooth Segmentation using MobileViT")
 
-    overlayed = cv2.addWeighted(image, 1.0, mask_color, 0.5, 0)
-    return overlayed
-
-# === UI ===
-st.title("🦷 Tooth Segmentation with MobileViT")
-st.write("Upload an image to segment the **visible teeth area** using a lightweight MobileViT segmentation model.")
-
-uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Upload a mouth image with visible teeth", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-    image = Image.open(uploaded_file)
-    tensor = preprocess_image(image)
+    image = Image.open(uploaded_file).convert("RGB")
+    input_tensor = transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        pred_mask = model(input_tensor)[0, 0].numpy()
 
-    with st.spinner("Segmenting..."):
-        with torch.no_grad():
-            pred = model(tensor)[0]
+    # Post-processing
+    pred_mask = (pred_mask > 0.7).astype(np.uint8) * 255
+    pred_mask = Image.fromarray(pred_mask).resize(image.size)
 
-    overlayed_img = overlay_mask(image, pred)
+    # Create overlay
+    overlay = Image.new("RGBA", image.size, (0, 0, 255, 100))  # Blue translucent
+    base = image.convert("RGBA")
+    pred_mask_rgba = Image.new("L", image.size, 0)
+    pred_mask_rgba.paste(255, mask=pred_mask)
+    final = Image.composite(overlay, base, pred_mask_rgba)
 
+    # Side-by-side display
     col1, col2 = st.columns(2)
     with col1:
         st.image(image, caption="Original Image", use_container_width=True)
     with col2:
-        st.image(overlayed_img, caption="Tooth Mask Overlay", use_container_width=True)
+        st.image(final, caption="Tooth Segmentation Overlay", use_container_width=True)