Update app.py

app.py

@@ -1,100 +1,122 @@
-import streamlit as st 
-
-import tensorflow as tf
-from PIL import Image
-import numpy as np
-import cv2
-from huggingface_hub import from_pretrained_keras
-
-
-try:
-    model=from_pretrained_keras("SerdarHelli/Segmentation-of-Teeth-in-Panoramic-X-ray-Image-Using-U-Net")
-except:
-    model=tf.keras.models.load_model("dental_xray_seg.h5")
-    pass
-    
-st.header("Segmentation of Teeth in Panoramic X-ray Image Using UNet")
-
-examples=["107.png","108.png","109.png"]
-link='Check Out Our Github Repo ! [link](https://github.com/SerdarHelli/Segmentation-of-Teeth-in-Panoramic-X-ray-Image-Using-U-Net)'
-st.markdown(link,unsafe_allow_html=True)
-
-
-def load_image(image_file):
-	img = Image.open(image_file)
-	return img
-
-def convert_one_channel(img):
-    #some images have 3 channels , although they are grayscale image
-    if len(img.shape)>2:
-        img= cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        return img
-    else:
-        return img
-    
-def convert_rgb(img):
-    #some images have 3 channels , although they are grayscale image
-    if len(img.shape)==2:
-        img= cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)  
-        return img
-    else:
-        return img
-    
-    
-st.subheader("Upload Dental Panoramic X-ray Image Image")
-image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])
-
-
-col1, col2, col3 = st.columns(3)
-with col1:
-    ex=load_image(examples[0])
-    st.image(ex,width=200)
-    if st.button('Example 1'):
-        image_file=examples[0]
-
-with col2:
-    ex1=load_image(examples[1])
-    st.image(ex1,width=200)
-    if st.button('Example 2'):
-        image_file=examples[1]
-
-
-with col3:
-    ex2=load_image(examples[2])
-    st.image(ex2,width=200)
-    if st.button('Example 3'):
-        image_file=examples[2]
-    
-    
-if image_file is not None:
-
-      img=load_image(image_file)
-      
-      st.text("Making A Prediction ....")
-      st.image(img,width=850)
-      
-      img=np.asarray(img)
-  
-      img_cv=convert_one_channel(img)
-      img_cv=cv2.resize(img_cv,(512,512), interpolation=cv2.INTER_LANCZOS4)
-      img_cv=np.float32(img_cv/255)
-      
-      img_cv=np.reshape(img_cv,(1,512,512,1))
-      prediction=model.predict(img_cv)
-      predicted=prediction[0]
-      predicted = cv2.resize(predicted, (img.shape[1],img.shape[0]), interpolation=cv2.INTER_LANCZOS4)
-      mask=np.uint8(predicted*255)# 
-      _, mask = cv2.threshold(mask, thresh=0, maxval=255, type=cv2.THRESH_BINARY+cv2.THRESH_OTSU)
-      kernel =( np.ones((5,5), dtype=np.float32))
-      mask=cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel,iterations=1 )  
-      mask=cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel,iterations=1 )
-      cnts,hieararch=cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
-      output = cv2.drawContours(convert_rgb(img), cnts, -1, (255, 0, 0) , 3)
-
-
-      if output is not None :      
-          st.subheader("Predicted Image")  
-          st.write(output.shape)
-          st.image(output,width=850)
-
-      st.text("DONE ! ....")
+import streamlit as st
+import tensorflow as tf
+from PIL import Image
+import numpy as np
+import cv2
+from huggingface_hub import from_pretrained_keras
+
+# Use st.cache_resource to load the model only once, preventing memory errors.
+@st.cache_resource
+def load_keras_model():
+    """Load the pre-trained Keras model from Hugging Face Hub and cache it."""
+    try:
+        # The model will be downloaded from the Hub and cached.
+        model = from_pretrained_keras("SerdarHelli/Segmentation-of-Teeth-in-Panoramic-X-ray-Image-Using-U-Net")
+        return model
+    except Exception as e:
+        # If model loading fails, show an error and return None.
+        st.error(f"Error loading the model: {e}")
+        return None
+
+# --- Helper Functions ---
+def load_image(image_file):
+    """Loads an image from a file path or uploaded file object."""
+    img = Image.open(image_file)
+    return img
+
+def convert_one_channel(img_array):
+    """Ensure the image is single-channel (grayscale)."""
+    # If image has 3 channels (like BGR or RGB), convert to grayscale.
+    if len(img_array.shape) > 2 and img_array.shape[2] > 1:
+        img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
+    return img_array
+
+def convert_rgb(img_array):
+    """Ensure the image is 3-channel (RGB) for drawing contours."""
+    # If image is grayscale, convert to RGB to draw colored contours.
+    if len(img_array.shape) == 2:
+        img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
+    return img_array
+
+# --- Streamlit App Layout ---
+st.header("Segmentation of Teeth in Panoramic X-ray Image Using UNet")
+
+link = 'Check Out Our Github Repo! [link](https://github.com/SerdarHelli/Segmentation-of-Teeth-in-Panoramic-X-ray-Image-Using-U-Net)'
+st.markdown(link, unsafe_allow_html=True)
+
+# Load the model and stop the app if it fails
+model = load_keras_model()
+if model is None:
+    st.warning("Model could not be loaded. The application cannot proceed.")
+    st.stop()
+
+# --- Image Selection Section ---
+st.subheader("Upload a Dental Panoramic X-ray Image or Select an Example")
+image_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
+
+st.write("---")
+st.write("Or choose an example:")
+examples = ["107.png", "108.png", "109.png"]
+col1, col2, col3 = st.columns(3)
+
+# Display example images and buttons to use them
+with col1:
+    st.image(examples[0], caption='Example 1', use_column_width=True)
+    if st.button('Use Example 1'):
+        image_file = examples[0]
+
+with col2:
+    st.image(examples[1], caption='Example 2', use_column_width=True)
+    if st.button('Use Example 2'):
+        image_file = examples[1]
+
+with col3:
+    st.image(examples[2], caption='Example 3', use_column_width=True)
+    if st.button('Use Example 3'):
+        image_file = examples[2]
+
+# --- Processing and Prediction Section ---
+if image_file is not None:
+    st.write("---")
+    
+    # Load and display the selected image
+    original_pil_img = load_image(image_file)
+    st.image(original_pil_img, caption="Original Image", use_column_width=True)
+    
+    with st.spinner("Analyzing image and predicting segmentation..."):
+        # Convert PIL image to NumPy array for processing
+        original_np_img = np.array(original_pil_img)
+
+        # 1. Pre-process for the model
+        img_gray = convert_one_channel(original_np_img.copy())
+        img_resized = cv2.resize(img_gray, (512, 512), interpolation=cv2.INTER_LANCZOS4)
+        img_normalized = np.float32(img_resized / 255.0)
+        img_input = np.reshape(img_normalized, (1, 512, 512, 1))
+
+        # 2. Make prediction
+        prediction = model.predict(img_input)
+        
+        # 3. Post-process the prediction mask
+        predicted_mask = prediction[0]
+        resized_mask = cv2.resize(predicted_mask, (original_np_img.shape[1], original_np_img.shape[0]), interpolation=cv2.INTER_LANCZOS4)
+        
+        # Binarize the mask using Otsu's thresholding
+        mask_8bit = (resized_mask * 255).astype(np.uint8)
+        _, final_mask = cv2.threshold(mask_8bit, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        
+        # Clean up mask with morphological operations
+        kernel = np.ones((5, 5), dtype=np.uint8)
+        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_OPEN, kernel, iterations=1)
+        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_CLOSE, kernel, iterations=1)
+        
+        # Find contours on the final mask
+        contours, _ = cv2.findContours(final_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Draw contours on a color version of the original image
+        img_for_drawing = convert_rgb(original_np_img.copy())
+        output_image = cv2.drawContours(img_for_drawing, contours, -1, (255, 0, 0), 3) # Draw red contours
+
+        st.subheader("Predicted Segmentation")
+        st.image(output_image, caption="Image with Segmented Teeth", use_column_width=True)
+        
+    st.success("Prediction complete!")