Update AKSHAYRAJAA/inference.py

AKSHAYRAJAA/inference.py

@@ -1,42 +1,40 @@
 import os
 import torch
-import spacy
-import spacy.cli
-from PIL import Image
-import torchvision.transforms as transforms
+import config
 import streamlit as st
+import spacy
+spacy.cli.download("en_core_web_sm")
+
 from utils import (
     load_dataset,
     get_model_instance,
     load_checkpoint,
     can_load_checkpoint,
+    normalize_text,
 )
-
-# Ensure the SpaCy model is downloaded
-spacy.cli.download("en_core_web_sm")
-nlp = spacy.load("en_core_web_sm")
+from PIL import Image
+import torchvision.transforms as transforms
 
 # Define device
 DEVICE = 'cpu'
 
 # Define image transformations
 TRANSFORMS = transforms.Compose([
-    transforms.Resize((224, 224)),  # Adjust to your model's input size
+    transforms.Resize((224, 224)),  # Replace with your model's expected input size
     transforms.ToTensor(),
     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()
-    vocabulary = dataset.vocab
+    dataset = load_dataset()  # Load dataset to access vocabulary
+    vocabulary = dataset.vocab  # Assuming 'vocab' is an attribute of the dataset
 
     st.write("Initializing the model...")
-    model = get_model_instance(vocabulary)
+    model = get_model_instance(vocabulary)  # Initialize the model
 
     if can_load_checkpoint():
         st.write("Loading checkpoint...")
@@ -48,58 +46,58 @@ 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")
-    image = TRANSFORMS(image).unsqueeze(0)
+    image = Image.open(image_path).convert("RGB")  # Ensure RGB format
+    image = TRANSFORMS(image).unsqueeze(0)  # Add batch dimension
     return image.to(DEVICE)
 
-
-def generate_report(model, image):
+def generate_report(model, image_path):
     """
     Generates a report for a given image using the model.
     """
-    st.write("Generating 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"])
-
-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()
-
-    # Preprocess and generate the report
     image = preprocess_image(image_path)
-    report = generate_report(model, image)
 
-    # Display the image and the report
-    if report:
-        st.image(image_path, caption="Uploaded Image", use_column_width=True)
-        st.write("Generated Report:")
+    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"])
+
+    if uploaded_file is not None:
+        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+        st.write("")
+
+        # 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.")
+
+        # 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)
+
+if __name__ == "__main__":
+    main()