带有 KerasHub 的 RAG 管道

import os
import sys

os.environ["KERAS_BACKEND"] = "jax"
import keras
import numpy as np

keras.config.set_dtype_policy("bfloat16")
import keras_hub
import tensorflow as tf
from PIL import Image
import matplotlib.pyplot as plt
from nilearn import datasets, image
import re

def load_models():
    """
    Load and configure vision model for feature extraction, Gemma3 VLM for report generation, and a compact text model for benchmarking.
    Returns:
        tuple: (vision_model, vlm_model, text_model)
    """
    # Vision model for feature extraction (lightweight MobileNetV3)
    vision_model = keras_hub.models.ImageClassifier.from_preset(
        "mobilenet_v3_large_100_imagenet_21k"
    )
    # Gemma3 Text model for report generation in RAG Pipeline (compact)
    text_model = keras_hub.models.Gemma3CausalLM.from_preset("gemma3_instruct_1b")
    # Gemma3 VLM for report generation (original, for benchmarking)
    vlm_model = keras_hub.models.Gemma3CausalLM.from_preset("gemma3_instruct_4b")
    return vision_model, vlm_model, text_model


# Load models
print("Loading models...")
vision_model, vlm_model, text_model = load_models()

def prepare_images_and_captions(oasis, images_dir="images"):
    """
    Prepare OASIS brain MRI images and generate captions.

    Args:
        oasis: OASIS dataset object containing brain MRI data
        images_dir (str): Directory to save processed images

    Returns:
        tuple: (image_paths, captions) - Lists of image paths and corresponding captions
    """
    os.makedirs(images_dir, exist_ok=True)
    image_paths = []
    captions = []
    for i, img_path in enumerate(oasis.gray_matter_maps):
        img = image.load_img(img_path)
        data = img.get_fdata()
        slice_ = data[:, :, data.shape[2] // 2]
        slice_ = (
            (slice_ - np.min(slice_)) / (np.max(slice_) - np.min(slice_)) * 255
        ).astype(np.uint8)
        img_pil = Image.fromarray(slice_)
        fname = f"oasis_{i}.png"
        fpath = os.path.join(images_dir, fname)
        img_pil.save(fpath)
        image_paths.append(fpath)
        captions.append(f"OASIS Brain MRI {i}")
    print("Saved 4 OASIS Brain MRI images:", image_paths)
    return image_paths, captions


# Prepare data
print("Preparing OASIS dataset...")
oasis = datasets.fetch_oasis_vbm(n_subjects=4)  # Use 4 images
print("Download dataset is completed.")
image_paths, captions = prepare_images_and_captions(oasis)

def visualize_images(image_paths, captions):
    """
    Visualize the processed brain MRI images.

    Args:
        image_paths (list): List of image file paths
        captions (list): List of corresponding image captions
    """
    n = len(image_paths)
    fig, axes = plt.subplots(1, n, figsize=(4 * n, 4))
    # If only one image, axes is not a list
    if n == 1:
        axes = [axes]
    for i, (img_path, title) in enumerate(zip(image_paths, captions)):
        img = Image.open(img_path)
        axes[i].imshow(img, cmap="gray")
        axes[i].set_title(title)
        axes[i].axis("off")
    plt.suptitle("OASIS Brain MRI Images")
    plt.tight_layout()
    plt.show()


# Visualize the prepared images
visualize_images(image_paths, captions)

def visualize_prediction(query_img_path, db_image_paths, best_idx, db_reports):
    """
    Visualize the query image and the most similar retrieved image.

    Args:
        query_img_path (str): Path to the query image
        db_image_paths (list): List of database image paths
        best_idx (int): Index of the most similar database image
        db_reports (list): List of database reports
    """
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    axes[0].imshow(Image.open(query_img_path), cmap="gray")
    axes[0].set_title("Query Image")
    axes[0].axis("off")
    axes[1].imshow(Image.open(db_image_paths[best_idx]), cmap="gray")
    axes[1].set_title("Retrieved Context Image")
    axes[1].axis("off")
    plt.suptitle("Query and Most Similar Database Image")
    plt.tight_layout()
    plt.show()

def extract_image_features(img_path, vision_model):
    """
    Extract features from an image using the vision model.

    Args:
        img_path (str): Path to the input image
        vision_model: Pre-trained vision model for feature extraction

    Returns:
        numpy.ndarray: Extracted feature vector
    """
    img = Image.open(img_path).convert("RGB").resize((384, 384))
    x = np.array(img) / 255.0
    x = np.expand_dims(x, axis=0)
    features = vision_model(x)
    return features

db_reports = [
    "MRI shows a 1.5cm lesion in the right frontal lobe, non-enhancing, no edema.",
    "Normal MRI scan, no abnormal findings.",
    "Diffuse atrophy noted, no focal lesions.",
]

def clean_generated_output(generated_text, prompt):
    """
    Remove prompt echo and header details from generated text.

    Args:
        generated_text (str): Raw generated text from the language model
        prompt (str): Original prompt used for generation

    Returns:
        str: Cleaned text without prompt echo and headers
    """
    # Remove the prompt from the beginning of the generated text
    if generated_text.startswith(prompt):
        cleaned_text = generated_text[len(prompt) :].strip()
    else:
        cleaned_text = generated_text.replace(prompt, "").strip()

    # Remove header details and unwanted formatting
    lines = cleaned_text.split("\n")
    filtered_lines = []
    skip_next = False
    subheading_pattern = re.compile(r"^(\s*[A-Za-z0-9 .\-()]+:)(.*)")

    for line in lines:
        line = line.replace("<end_of_turn>", "").strip()
        line = line.replace("**", "")
        line = line.replace("*", "")
        # Remove empty lines after headers (existing logic)
        if any(
            header in line
            for header in [
                "**Patient:**",
                "**Date of Exam:**",
                "**Exam:**",
                "**Referring Physician:**",
                "**Patient ID:**",
                "Patient:",
                "Date of Exam:",
                "Exam:",
                "Referring Physician:",
                "Patient ID:",
            ]
        ):
            continue
        elif line.strip() == "" and skip_next:
            skip_next = False
            continue
        else:
            # Split subheadings onto their own line if content follows
            match = subheading_pattern.match(line)
            if match and match.group(2).strip():
                filtered_lines.append(match.group(1).strip())
                filtered_lines.append(match.group(2).strip())
                filtered_lines.append("")  # Add a blank line after subheading
            else:
                filtered_lines.append(line)
                # Add a blank line after subheadings (lines ending with ':')
                if line.endswith(":") and (
                    len(filtered_lines) == 1 or filtered_lines[-2] != ""
                ):
                    filtered_lines.append("")
            skip_next = False

    # Remove any empty lines and excessive whitespace
    cleaned_text = "\n".join(
        [l for l in filtered_lines if l.strip() or l == ""]
    ).strip()

    return cleaned_text

def rag_pipeline(query_img_path, db_image_paths, db_reports, vision_model, text_model):
    """
    Retrieval-Augmented Generation pipeline using vision model for retrieval and a compact text model for report generation.
    Args:
        query_img_path (str): Path to the query image
        db_image_paths (list): List of database image paths
        db_reports (list): List of database reports
        vision_model: Vision model for feature extraction
        text_model: Compact text model for report generation
    Returns:
        tuple: (best_idx, retrieved_report, generated_report)
    """
    # Extract features for the query image
    query_features = extract_image_features(query_img_path, vision_model)
    # Extract features for the database images
    db_features = np.vstack(
        [extract_image_features(p, vision_model) for p in db_image_paths]
    )
    # Ensure features are numpy arrays for similarity search
    db_features_np = np.array(db_features)
    query_features_np = np.array(query_features)
    # Similarity search
    similarity = np.dot(db_features_np, query_features_np.T).squeeze()
    best_idx = np.argmax(similarity)
    retrieved_report = db_reports[best_idx]
    print(f"[RAG] Matched image index: {best_idx}")
    print(f"[RAG] Matched image path: {db_image_paths[best_idx]}")
    print(f"[RAG] Retrieved context/report:\n{retrieved_report}\n")
    PROMPT_TEMPLATE = (
        "Context:\n{context}\n\n"
        "Based on the above radiology report and the provided brain MRI image, please:\n"
        "1. Provide a diagnostic impression.\n"
        "2. Explain the diagnostic reasoning.\n"
        "3. Suggest possible treatment options.\n"
        "Format your answer as a structured radiology report.\n"
    )
    prompt = PROMPT_TEMPLATE.format(context=retrieved_report)
    # Generate report using the text model (text only, no image input)
    output = text_model.generate(
        {
            "prompts": prompt,
        }
    )
    cleaned_output = clean_generated_output(output, prompt)
    return best_idx, retrieved_report, cleaned_output


# Split data: first 3 as database, last as query
db_image_paths = image_paths[:-1]
query_img_path = image_paths[-1]

# Run RAG pipeline
print("Running RAG pipeline...")
best_idx, retrieved_report, generated_report = rag_pipeline(
    query_img_path, db_image_paths, db_reports, vision_model, text_model
)

# Visualize results
visualize_prediction(query_img_path, db_image_paths, best_idx, db_reports)

# Print RAG results
print("\n" + "=" * 50)
print("RAG PIPELINE RESULTS")
print("=" * 50)
print(f"\nMatched DB Report Index: {best_idx}")
print(f"Matched DB Report: {retrieved_report}")
print("\n--- Generated Report ---\n", generated_report)

def vlm_generate_report(query_img_path, vlm_model, question=None):
    """
    Generate a radiology report directly from the image using a vision-language model.
    Args:
        query_img_path (str): Path to the query image
        vlm_model: Pre-trained vision-language model (Gemma3 4B VLM)
        question (str): Optional question or prompt to include
    Returns:
        str: Generated radiology report
    """
    PROMPT_TEMPLATE = (
        "Based on the provided brain MRI image, please:\n"
        "1. Provide a diagnostic impression.\n"
        "2. Explain the diagnostic reasoning.\n"
        "3. Suggest possible treatment options.\n"
        "Format your answer as a structured radiology report.\n"
    )
    if question is None:
        question = ""
    # Preprocess the image as required by the model
    img = Image.open(query_img_path).convert("RGB").resize((224, 224))
    image = np.array(img) / 255.0
    image = np.expand_dims(image, axis=0)
    # Generate report using the VLM
    output = vlm_model.generate(
        {
            "images": image,
            "prompts": PROMPT_TEMPLATE.format(question=question),
        }
    )
    # Clean the generated output
    cleaned_output = clean_generated_output(
        output, PROMPT_TEMPLATE.format(question=question)
    )
    return cleaned_output


# Run VLM (direct approach)
print("\n" + "=" * 50)
print("VLM RESULTS (Direct Approach)")
print("=" * 50)
vlm_report = vlm_generate_report(query_img_path, vlm_model)
print("\n--- Vision-Language Model (No Retrieval) Report ---\n", vlm_report)