Update inference.py

inference.py

@@ -1,9 +1,7 @@
 import os
 import torch
-import config
 import spacy
-spacy.cli.download("en_core_web_sm")
-
+import config
 from utils import (
     load_dataset,
     get_model_instance,
@@ -15,6 +13,9 @@ from PIL import Image
 import torchvision.transforms as transforms
 import streamlit as st
 
+# Download Spacy model (only once during runtime)
+spacy.cli.download("en_core_web_sm")
+
 # Define device
 DEVICE = 'cpu'
 
@@ -26,6 +27,7 @@ TRANSFORMS = transforms.Compose([
 ])
 
 
+@st.cache_resource
 def load_model():
     """
     Loads the model with the vocabulary and checkpoint.
@@ -39,9 +41,13 @@ def load_model():
 
     if can_load_checkpoint():
         st.write("Loading checkpoint...")
-        load_checkpoint(model)
+        try:
+            load_checkpoint(model)
+        except RuntimeError as e:
+            st.error(f"Error loading checkpoint: {e}")
+            st.stop()
     else:
-        st.write("No checkpoint found, starting with untrained model.")
+        st.warning("No checkpoint found, starting with untrained model.")
 
     model.eval()  # Set the model to evaluation mode
     st.write("Model is ready for inference.")
@@ -53,9 +59,13 @@ def preprocess_image(image_path):
     Preprocess the input image for the model.
     """
     st.write(f"Preprocessing image: {image_path}")
-    image = Image.open(image_path).convert("RGB")
-    image = TRANSFORMS(image).unsqueeze(0)
-    return image.to(DEVICE)
+    try:
+        image = Image.open(image_path).convert("RGB")
+        image = TRANSFORMS(image).unsqueeze(0)
+        return image.to(DEVICE)
+    except Exception as e:
+        st.error(f"Error preprocessing image: {e}")
+        st.stop()
 
 
 def generate_report(model, image):
@@ -63,12 +73,15 @@ def generate_report(model, image):
     Generates a report for a given image using the model.
     """
     st.write("Generating report...")
-    with torch.no_grad():
-        output = model.generate_caption(image, max_length=25)
-        report = " ".join(output)
-
-    st.write(f"Generated report: {report}")
-    return report
+    try:
+        with torch.no_grad():
+            output = model.generate_caption(image, max_length=25)
+            report = " ".join(output)
+        st.write(f"Generated report: {report}")
+        return report
+    except Exception as e:
+        st.error(f"Error generating report: {e}")
+        st.stop()
 
 
 # Streamlit App
@@ -79,6 +92,9 @@ st.write("Upload an X-ray image to generate a report.")
 uploaded_file = st.file_uploader("Choose an image file", type=["png", "jpg", "jpeg"])
 
 if uploaded_file is not None:
+    # Ensure 'temp' directory exists
+    os.makedirs("temp", exist_ok=True)
+
     # Save uploaded file to disk
     image_path = os.path.join("temp", uploaded_file.name)
     with open(image_path, "wb") as f:
@@ -94,4 +110,4 @@ if uploaded_file is not None:
     # Display the image and the report
     st.image(image_path, caption="Uploaded Image", use_column_width=True)
     st.write("Generated Report:")
-    st.write(report)
+    st.write(report)