Update AKSHAYRAJAA/inference.py

AKSHAYRAJAA/inference.py

@@ -1,18 +1,18 @@
 import os
 import torch
-import config
 import spacy
 import spacy.cli
+from PIL import Image
+import torchvision.transforms as transforms
+import streamlit as st
 from utils import (
     load_dataset,
     get_model_instance,
     load_checkpoint,
     can_load_checkpoint,
-    normalize_text,
 )
-from PIL import Image
-import torchvision.transforms as transforms
-import streamlit as st
+
+# Ensure the SpaCy model is downloaded
 spacy.cli.download("en_core_web_sm")
 nlp = spacy.load("en_core_web_sm")
 
@@ -21,7 +21,7 @@ DEVICE = 'cpu'
 
 # Define image transformations
 TRANSFORMS = transforms.Compose([
-    transforms.Resize((224, 224)),  # Replace with your model's expected input size
+    transforms.Resize((224, 224)),  # Adjust to your model's input size
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
@@ -64,18 +64,24 @@ 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)  # Ensure `generate_caption` is implemented
+            report = " ".join(output)
+        st.write(f"Generated report: {report}")
+        return report
+    except Exception as e:
+        st.error(f"Error during report generation: {e}")
+        return None
 
 
 # Streamlit App
 st.title("Medical Image Report Generator")
 st.write("Upload an X-ray image to generate a report.")
 
+# Create temp directory
+os.makedirs("temp", exist_ok=True)
+
 # File uploader
 uploaded_file = st.file_uploader("Choose an image file", type=["png", "jpg", "jpeg"])
 
@@ -93,6 +99,7 @@ if uploaded_file is not None:
     report = generate_report(model, image)
 
     # Display the image and the report
-    st.image(image_path, caption="Uploaded Image", use_column_width=True)
-    st.write("Generated Report:")
-    st.write(report)
+    if report:
+        st.image(image_path, caption="Uploaded Image", use_column_width=True)
+        st.write("Generated Report:")
+        st.write(report)