Update app.py

app.py

@@ -1,27 +1,22 @@
 import streamlit as st
 import torch
 import torch.nn as nn
+import torchvision.transforms as transforms
 import timm
 import numpy as np
 import cv2
 from PIL import Image
 import io
+import warnings
 
-# Hide Streamlit warnings and UI elements
-st.set_page_config(layout="wide")
-st.markdown("""
-    <style>
-        footer {visibility: hidden;}
-    </style>
-""", unsafe_allow_html=True)
+warnings.filterwarnings("ignore")
 
-# === Model Definition ===
+# Define the model class
 class MobileViTSegmentation(nn.Module):
     def __init__(self, encoder_name='mobilevit_s', pretrained=False):
         super().__init__()
         self.backbone = timm.create_model(encoder_name, features_only=True, pretrained=pretrained)
         self.encoder_channels = self.backbone.feature_info.channels()
-
         self.decoder = nn.Sequential(
             nn.Conv2d(self.encoder_channels[-1], 128, kernel_size=3, padding=1),
             nn.Upsample(scale_factor=2, mode='bilinear'),
@@ -39,55 +34,53 @@ 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 function
 @st.cache_resource
 def load_model():
     model = MobileViTSegmentation()
-    state_dict = torch.load("mobilevit_teeth_segmentation.pth", map_location="cpu")
-    model.load_state_dict(state_dict)
+    model.load_state_dict(torch.load("mobilevit_teeth_segmentation.pth", 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
-
-    mask_color = np.zeros_like(image)
-    mask_color[..., 2] = mask_bin  # Blue mask
-
-    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"])
+# Inference
+def predict_mask(image, model, threshold=0.7):
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor()
+    ])
+    img_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        pred = model(img_tensor)
+    pred_mask = pred.squeeze().numpy()
+    pred_mask = (pred_mask > threshold).astype(np.uint8)
+    return pred_mask
+
+# Overlay mask on image
+def overlay_mask(image, mask, color=(0, 0, 255), alpha=0.4):
+    image_np = np.array(image.convert("RGB"))
+    mask_resized = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]))
+    color_mask = np.zeros_like(image_np)
+    color_mask[:, :] = color
+    overlay = np.where(mask_resized[..., None] == 1, color_mask, 0)
+    blended = cv2.addWeighted(image_np, 1 - alpha, overlay, alpha, 0)
+    return blended
+
+# Streamlit UI
+st.title("🦷 Tooth Segmentation from Mouth Images")
+st.markdown("Upload a face or mouth image and get the segmented **tooth region overlayed**.")
+
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-    image = Image.open(uploaded_file)
-    tensor = preprocess_image(image)
+    image = Image.open(uploaded_file).convert("RGB")
+    model = load_model()
+    pred_mask = predict_mask(image, model)
 
-    with st.spinner("Segmenting..."):
-        with torch.no_grad():
-            pred = model(tensor)[0]
-
-    overlayed_img = overlay_mask(image, pred)
+    overlayed_img = overlay_mask(image, pred_mask, color=(0, 0, 255), alpha=0.4)
 
     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)
+