Update AKSHAYRAJAA/inference.py

AKSHAYRAJAA/inference.py

@@ -10,11 +10,12 @@ from utils import (
 )
 from PIL import Image
 import torchvision.transforms as transforms
+import streamlit as st
 
 # 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(),
@@ -26,21 +27,21 @@ def load_model():
     """
     Loads the model with the vocabulary and checkpoint.
     """
-    print("Loading dataset and vocabulary...")
-    dataset = load_dataset()  # Load dataset to access vocabulary
-    vocabulary = dataset.vocab  # Assuming 'vocab' is an attribute of the dataset
+    st.write("Loading dataset and vocabulary...")
+    dataset = load_dataset()
+    vocabulary = dataset.vocab
 
-    print("Initializing the model...")
-    model = get_model_instance(vocabulary)  # Initialize the model
+    st.write("Initializing the model...")
+    model = get_model_instance(vocabulary)
 
     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
 
 
@@ -48,41 +49,46 @@ def preprocess_image(image_path):
     """
     Preprocess the input image for the model.
     """
-    print(f"Preprocessing image: {image_path}")
-    image = Image.open(image_path).convert("RGB")  # Ensure RGB format
-    image = TRANSFORMS(image).unsqueeze(0)  # Add batch dimension
+    st.write(f"Preprocessing image: {image_path}")
+    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)
-
-    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
 
 
-if __name__ == "__main__":
-    # Path to the checkpoint file
-    CHECKPOINT_PATH = config.CHECKPOINT_FILE  # Ensure config.CHECKPOINT_FILE is correctly set
+# Streamlit App
+st.title("Medical Image Report Generator")
+st.write("Upload an X-ray image to generate a report.")
+
+# File uploader
+uploaded_file = st.file_uploader("Choose an image file", type=["png", "jpg", "jpeg"])
 
-    # Path to the input image
-    IMAGE_PATH = "D:\AKSHAYRAJAA\dataset\images\CXR387_IM-1962-1001.png"  # Replace with your image path
+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()
 
-    # 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}")
+    # Preprocess and generate the report
+    image = preprocess_image(image_path)
+    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)