Update inference.py

inference.py

@@ -1,10 +1,10 @@
 import os
 import torch
 import config
-import streamlit as st
-import spacy
-spacy.cli.download("en_core_web_sm")
+!python -m spacy download en_core_web_sm
 
+import spacy
+nlp = spacy.load("en_core_web_sm")
 from utils import (
     load_dataset,
     get_model_instance,
@@ -14,6 +14,7 @@ from utils import (
 )
 from PIL import Image
 import torchvision.transforms as transforms
+import streamlit as st
 
 # Define device
 DEVICE = 'cpu'
@@ -25,16 +26,17 @@ TRANSFORMS = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
+
 def load_model():
     """
     Loads the model with the vocabulary and checkpoint.
     """
     st.write("Loading dataset and vocabulary...")
-    dataset = load_dataset()  # Load dataset to access vocabulary
-    vocabulary = dataset.vocab  # Assuming 'vocab' is an attribute of the dataset
+    dataset = load_dataset()
+    vocabulary = dataset.vocab
 
     st.write("Initializing the model...")
-    model = get_model_instance(vocabulary)  # Initialize the model
+    model = get_model_instance(vocabulary)
 
     if can_load_checkpoint():
         st.write("Loading checkpoint...")
@@ -46,58 +48,51 @@ def load_model():
     st.write("Model is ready for inference.")
     return model
 
+
 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")  # Ensure RGB format
-    image = TRANSFORMS(image).unsqueeze(0)  # Add batch dimension
+    image = Image.open(image_path).convert("RGB")
+    image = TRANSFORMS(image).unsqueeze(0)
     return image.to(DEVICE)
 
-def generate_report(model, image_path):
+
+def generate_report(model, image):
     """
     Generates a report for a given image using the model.
     """
-    image = preprocess_image(image_path)
-
     st.write("Generating report...")
     with torch.no_grad():
-        # Assuming the model has a 'generate_caption' method
         output = model.generate_caption(image, max_length=25)
         report = " ".join(output)
 
     st.write(f"Generated report: {report}")
     return report
 
-# Streamlit app
-def main():
-    st.title("Chest X-Ray Report Generator")
-    st.write("Upload a Chest X-Ray image to generate a medical report.")
 
-    # Upload image
-    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+# Streamlit App
+st.title("Medical Image Report Generator")
+st.write("Upload an X-ray image to generate a report.")
 
-    if uploaded_file is not None:
-        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
-        st.write("")
+# File uploader
+uploaded_file = st.file_uploader("Choose an image file", type=["png", "jpg", "jpeg"])
 
-        # Save the uploaded file temporarily
-        image_path = "./temp_image.png"
-        with open(image_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
-        st.write("Image uploaded successfully.")
+if uploaded_file is not None:
+    # Save uploaded file to disk
+    image_path = os.path.join("temp", uploaded_file.name)
+    with open(image_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
 
-        # Load the model
-        model = load_model()
+    # Load the model
+    model = load_model()
 
-        # Generate report
-        report = generate_report(model, image_path)
-        st.write("### Generated Report:")
-        st.write(report)
-
-        # Clean up temporary file
-        os.remove(image_path)
+    # Preprocess and generate the report
+    image = preprocess_image(image_path)
+    report = generate_report(model, image)
 
-if __name__ == "__main__":
-    main()
+    # Display the image and the report
+    st.image(image_path, caption="Uploaded Image", use_column_width=True)
+    st.write("Generated Report:")
+    st.write(report)