Spaces:
Sleeping
Sleeping
File size: 3,015 Bytes
c20fb92 774af3e c20fb92 2851f3a c20fb92 774af3e c20fb92 774af3e 2851f3a 774af3e c20fb92 774af3e c20fb92 774af3e c20fb92 774af3e c20fb92 774af3e c20fb92 774af3e c20fb92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import streamlit as st
from PIL import Image
from transformers import pipeline
import io
import torch # Import PyTorch
# --- Configuration ---
# Specify the model
MODEL_NAME = "microsoft/maira-2"
# --- Model Loading (using pipeline) ---
@st.cache_resource # Cache the pipeline for performance
def load_pipeline():
"""Loads the VQA pipeline."""
try:
# Explicitly set device if CUDA is available, otherwise use CPU
device = 0 if torch.cuda.is_available() else -1 # Use torch.cuda
vqa_pipeline = pipeline("visual-question-answering", model=MODEL_NAME, device=device) # Add device
return vqa_pipeline
except Exception as e:
st.error(f"Error loading pipeline: {e}")
return None
# --- Image Preprocessing (Keep as bytes) ---
def prepare_image(image):
"""Prepares the PIL Image object for the pipeline (handles RGBA)."""
image_bytes = io.BytesIO()
if image.mode == "RGBA":
image = image.convert("RGB")
image.save(image_bytes, format="JPEG")
return image_bytes.getvalue() # Return bytes directly
# --- Streamlit App ---
def main():
st.title("Chest X-ray Analysis with Maira-2 (Transformers Pipeline)")
st.write("Upload a chest X-ray image. This app uses the Maira-2 model via the Transformers library.")
vqa_pipeline = load_pipeline()
if vqa_pipeline is None:
st.warning("Pipeline not loaded. Predictions will not be available.")
return
uploaded_file = st.file_uploader("Choose a chest X-ray image (JPG, PNG)", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image", use_column_width=True)
with st.spinner("Analyzing image with Maira-2..."):
image_data = prepare_image(image)
try:
results = vqa_pipeline(
image=image_data, # Pass the image bytes
question="Analyze this chest X-ray image and provide detailed findings. Include any abnormalities, their locations, and potential diagnoses. Be as specific as possible.",
)
if results: # Handle results (list of dicts)
if isinstance(results, list) and len(results) > 0:
best_answer = max(results, key=lambda x: x.get('score', 0))
if 'answer' in best_answer:
st.subheader("Findings:")
st.write(best_answer['answer'])
else:
st.warning("Could not find 'answer' in results.")
else:
st.warning("Unexpected result format.")
except Exception as e:
st.error(f"An error occurred during analysis: {e}")
else:
st.write("Please upload an image.")
st.write("---")
st.write("Disclaimer: For informational purposes only. Not medical advice.")
if __name__ == "__main__":
main() |