Update src/streamlit_app.py

src/streamlit_app.py

@@ -4,12 +4,10 @@ import torch.nn as nn
 import timm
 import numpy as np
 from PIL import Image
-import requests
-from io import BytesIO
 import torchvision.transforms as T
-import matplotlib.pyplot as plt
 from huggingface_hub import hf_hub_download
 
+# ========== Model Definition ==========
 class MobileViTSegmentation(nn.Module):
     def __init__(self, encoder_name='mobilevit_s', pretrained=True):
         super().__init__()
@@ -36,13 +34,12 @@ class MobileViTSegmentation(nn.Module):
 # ========== Load Model ==========
 @st.cache_resource
 def load_model():
-    from model import MobileViTSegmentation  # or define it here if needed
-    cache_dir = "/tmp/huggingface"  # Writable path in HF Spaces
+    cache_dir = "/tmp/huggingface"  # Safe writable directory in HF Spaces
 
     checkpoint_path = hf_hub_download(
         repo_id="svsaurav95/ToothSegmentation",
         filename="mobilevit_teeth_segmentation.pth",
-        cache_dir=cache_dir  # <- this fixes the permission issue
+        cache_dir=cache_dir
     )
 
     model = MobileViTSegmentation(pretrained=False)
@@ -52,14 +49,15 @@ def load_model():
 
 model = load_model()
 
-# ========== Image Transformation ==========
+# ========== Image Preprocessing ==========
 transform = T.Compose([
     T.Resize((256, 256)),
     T.ToTensor()
 ])
 
-# ========== Streamlit UI ==========
-st.title("Tooth Segmentation using MobileViT")
+# ========== UI ==========
+st.set_page_config(page_title="Tooth Segmentation", layout="wide")
+st.title("🦷 Tooth Segmentation using MobileViT")
 
 uploaded_file = st.file_uploader("Upload a mouth image with visible teeth", type=["jpg", "jpeg", "png"])
 
@@ -70,12 +68,12 @@ if uploaded_file:
     with torch.no_grad():
         pred_mask = model(input_tensor)[0, 0].numpy()
 
-    # Post-processing
+    # Threshold and resize to original
     pred_mask = (pred_mask > 0.7).astype(np.uint8) * 255
     pred_mask = Image.fromarray(pred_mask).resize(image.size)
 
-    # Create overlay
-    overlay = Image.new("RGBA", image.size, (0, 0, 255, 100))  
+    # Create translucent blue overlay
+    overlay = Image.new("RGBA", image.size, (0, 0, 255, 100))
     base = image.convert("RGBA")
     pred_mask_rgba = Image.new("L", image.size, 0)
     pred_mask_rgba.paste(255, mask=pred_mask)
@@ -86,4 +84,4 @@ if uploaded_file:
     with col1:
         st.image(image, caption="Original Image", use_container_width=True)
     with col2:
-        st.image(final, caption="Tooth Segmentation Overlay", use_container_width=True)
+        st.image(final, caption="Tooth Area Segmentation", use_container_width=True)