Update image_pipeline.py

image_pipeline.py

@@ -1,23 +1,21 @@
-import streamlit as st
-from transformers import pipeline
-from config import IMAGE_MODEL_NAME
+from transformers import AutoProcessor, AutoModelForImageTextToText
+from PIL import Image
+import torch
 
-@st.cache_resource
-def load_image_model():
-    """
-    Loads an image captioning model recognized by the HF pipeline.
-    Example: "nlpconnect/vit-gpt2-image-captioning" or "Salesforce/blip-image-captioning-base".
-    """
-    return pipeline("image-to-text", model=IMAGE_MODEL_NAME)
+from config import HF_IMAGE_MODEL
+
+# Load the advanced vision-language model for medical images
+processor = AutoProcessor.from_pretrained(HF_IMAGE_MODEL)
+model = AutoModelForImageTextToText.from_pretrained(HF_IMAGE_MODEL)
 
-def analyze_image(image_file, image_model):
+def analyze_medical_image(image_file):
     """
-    Pass an uploaded image to the loaded pipeline for caption generation.
+    Performs advanced medical image analysis. 
+    Returns a text explanation or diagnostic insight from the model.
     """
-    try:
-        result = image_model(image_file)
-        if isinstance(result, list) and len(result) > 0:
-            return result[0].get("generated_text", "No caption.")
-        return "No output from the model."
-    except Exception as e:
-        return f"Error analyzing image: {str(e)}"
+    image = Image.open(image_file).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt").to(model.device)
+    
+    # Inference
+    outputs = model.generate(**inputs, max_length=256)
+    return processor.batch_decode(outputs, skip_special_tokens=True)[0]