Update inference.py

inference.py

@@ -1,6 +1,7 @@
 import os
 import torch
 import config
+import streamlit as st
 from utils import (
     load_dataset,
     get_model_instance,
@@ -14,75 +15,86 @@ import torchvision.transforms as transforms
 # Define device
 DEVICE = 'cpu'
 
-# Define image transformations (adjust based on training setup)
+# Define image transformations
 TRANSFORMS = transforms.Compose([
     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.
     """
-    print("Loading dataset and vocabulary...")
+    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
 
-    print("Initializing the model...")
+    st.write("Initializing the model...")
     model = get_model_instance(vocabulary)  # Initialize the model
 
     if can_load_checkpoint():
-        print("Loading checkpoint...")
+        st.write("Loading checkpoint...")
         load_checkpoint(model)
     else:
-        print("No checkpoint found, starting with untrained model.")
+        st.write("No checkpoint found, starting with untrained model.")
 
     model.eval()  # Set the model to evaluation mode
-    print("Model is ready for inference.")
+    st.write("Model is ready for inference.")
     return model
 
-
 def preprocess_image(image_path):
     """
     Preprocess the input image for the model.
     """
-    print(f"Preprocessing image: {image_path}")
+    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
     return image.to(DEVICE)
 
-
 def generate_report(model, image_path):
     """
     Generates a report for a given image using the model.
     """
     image = preprocess_image(image_path)
 
-    print("Generating 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)
 
-    print(f"Generated report: {report}")
+    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.")
 
-if __name__ == "__main__":
-    # Path to the checkpoint file
-    CHECKPOINT_PATH = config.CHECKPOINT_FILE  # Ensure config.CHECKPOINT_FILE is correctly set
+    # Upload image
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
-    # Path to the input image
-    IMAGE_PATH = "./dataset/images/CXR1178_IM-0121-1001.png"  # Replace with your image path
+    if uploaded_file is not None:
+        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+        st.write("")
 
-    # Load the model
-    model = load_model()
+        # 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.")
 
-    # Ensure the image exists before inference
-    if os.path.exists(IMAGE_PATH):
-        report = generate_report(model, IMAGE_PATH)
-        print("Final Report:", report)
-    else:
-        print(f"Image not found at path: {IMAGE_PATH}")
+        # 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()