FeatureX-Clinical-AI

Running on Zero

Upload 11 files
12d3559 verified 4 days ago
45.9 kB
import os
import io
import gc
import logging
import re
import threading
from datetime import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms

from PIL import Image, ImageOps, UnidentifiedImageError
import numpy as np
import cv2

from janus_engine import generate_reports as janus_generate_reports, parse_radiology_text
from arabic_translation import translate_report_en_to_ar
from report_quality import check_bilingual_reports


# ============================================================
# Clinical AI X-Ray - Inference Core
# EfficientNet-B4 + Grad-CAM
# ============================================================

os.environ.setdefault("OMP_NUM_THREADS", "4")
os.environ.setdefault("MKL_NUM_THREADS", "4")

try:
    torch.set_num_threads(4)
except Exception:
    pass

try:
    if hasattr(torch, "set_num_interop_threads"):
        torch.set_num_interop_threads(1)
except Exception:
    pass


logger = logging.getLogger("clinical-ai")

MODEL_PATH = os.environ.get(
    "MODEL_PATH",
    os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "best_efficientnet.pth",
    ),
)

IMG_SIZE = 380
MAX_FILE_SIZE = 15 * 1024 * 1024

MIN_WIDTH = 200
MIN_HEIGHT = 200
MAX_WIDTH = 4000
MAX_HEIGHT = 4000

Image.MAX_IMAGE_PIXELS = 100_000_000

CLASSES = [
    "COVID-19",
    "Cardiomegaly",
    "Normal",
    "Pleural Effusion",
    "Pneumonia",
    "Tuberculosis",
]


REPORT_TEMPLATES_AR = {
    "COVID-19": {
        "title": "اشتباه إصابة بكوفيد-19",
        "findings": "توجد مظاهر رئوية قد تتوافق مع التهاب رئوي فيروسي.",
        "impression": "النتيجة الآلية تشير إلى احتمال وجود إصابة بكوفيد-19.",
        "recommendations": [
            "مراجعة الطبيب المختص.",
            "النظر في إجراء فحص PCR حسب التقييم السريري.",
            "تقييم الأعراض والعلامات الحيوية.",
        ],
        "differential_diagnosis": [
            "التهاب رئوي فيروسي آخر",
            "التهاب رئوي جرثومي",
            "وذمة رئوية",
        ],
        "urgency": "high",
        "follow_up": "حسب الحالة السريرية وتقييم الطبيب.",
    },
    "Cardiomegaly": {
        "title": "اشتباه تضخم القلب",
        "findings": "توجد مؤشرات شعاعية قد تتوافق مع زيادة حجم ظل القلب.",
        "impression": "النتيجة الآلية تشير إلى احتمال وجود تضخم في القلب.",
        "recommendations": [
            "مراجعة طبيب القلب.",
            "النظر في إجراء Echocardiography.",
            "ربط النتيجة بالفحص السريري.",
        ],
        "differential_diagnosis": [
            "انصباب تاموري",
            "اعتلال عضلة القلب",
            "تضخم بطيني",
        ],
        "urgency": "medium",
        "follow_up": "حسب تقييم الطبيب.",
    },
    "Normal": {
        "title": "صورة صدر تبدو طبيعية",
        "findings": "لم يحدد النموذج علامات واضحة تتوافق مع الفئات المرضية المدربة عليها.",
        "impression": "لم تظهر علامات مرضية واضحة وفق تصنيف النموذج.",
        "recommendations": [
            "ربط النتيجة بالحالة السريرية.",
            "مراجعة أخصائي الأشعة عند الحاجة.",
        ],
        "differential_diagnosis": [],
        "urgency": "low",
        "follow_up": "حسب الحاجة السريرية.",
    },
    "Pleural Effusion": {
        "title": "اشتباه انصباب جنبي",
        "findings": "توجد مؤشرات قد تتوافق مع وجود تجمع سائل في الحيز الجنبي.",
        "impression": "النتيجة الآلية تشير إلى احتمال وجود انصباب جنبي.",
        "recommendations": [
            "مراجعة طبيب صدرية أو باطنية.",
            "ربط النتيجة بالفحص السريري.",
            "النظر في التصوير الإضافي حسب قرار الطبيب.",
        ],
        "differential_diagnosis": [
            "قصور القلب",
            "التهاب رئوي",
            "أمراض جنبية أخرى",
        ],
        "urgency": "medium",
        "follow_up": "حسب الأعراض والتقييم السريري.",
    },
    "Pneumonia": {
        "title": "اشتباه التهاب رئوي",
        "findings": "توجد مؤشرات شعاعية قد تتوافق مع التهاب رئوي.",
        "impression": "النتيجة الآلية تشير إلى احتمال وجود التهاب رئوي.",
        "recommendations": [
            "مراجعة الطبيب.",
            "ربط الصورة بالأعراض والفحوصات المخبرية.",
            "تقييم مستوى الأكسجين عند الحاجة.",
        ],
        "differential_diagnosis": [
            "التهاب رئوي فيروسي",
            "التهاب رئوي جرثومي",
            "وذمة رئوية",
        ],
        "urgency": "high",
        "follow_up": "حسب شدة الأعراض وتقييم الطبيب.",
    },
    "Tuberculosis": {
        "title": "اشتباه سل رئوي",
        "findings": "توجد مؤشرات شعاعية قد تتوافق مع الإصابة بالسل الرئوي.",
        "impression": "النتيجة الآلية تشير إلى احتمال وجود سل رئوي.",
        "recommendations": [
            "مراجعة طبيب صدرية أو أمراض معدية.",
            "إجراء الفحوصات المناسبة لتأكيد التشخيص.",
            "عدم الاعتماد على نتيجة النموذج وحدها.",
        ],
        "differential_diagnosis": [
            "التهاب رئوي",
            "أمراض رئوية أخرى",
            "أورام رئوية",
        ],
        "urgency": "high",
        "follow_up": "متابعة عاجلة مع الطبيب المختص.",
    },
}


REPORT_TEMPLATES_EN = {
    "COVID-19": {
        "title": "Suspected COVID-19 Infection",
        "findings": "There are pulmonary findings that may be consistent with viral pneumonia.",
        "impression": "The automated result suggests a possible COVID-19 infection.",
        "recommendations": [
            "Consult a specialist physician.",
            "Consider PCR testing based on clinical evaluation.",
            "Assess symptoms and vital signs.",
        ],
        "differential_diagnosis": [
            "Other viral pneumonia",
            "Bacterial pneumonia",
            "Pulmonary edema",
        ],
        "urgency": "high",
        "follow_up": "Based on clinical condition and physician assessment.",
    },
    "Cardiomegaly": {
        "title": "Suspected Cardiomegaly",
        "findings": "Radiographic signs may be consistent with increased cardiac shadow size.",
        "impression": "The automated result suggests possible cardiomegaly.",
        "recommendations": [
            "Consult a cardiologist.",
            "Consider echocardiography.",
            "Correlate with clinical examination.",
        ],
        "differential_diagnosis": [
            "Pericardial effusion",
            "Cardiomyopathy",
            "Ventricular hypertrophy",
        ],
        "urgency": "medium",
        "follow_up": "Based on physician evaluation.",
    },
    "Normal": {
        "title": "Chest X-ray Appears Normal",
        "findings": "The model did not identify clear signs matching the trained disease categories.",
        "impression": "No obvious pathological signs were detected according to the model classification.",
        "recommendations": [
            "Correlate with clinical condition.",
            "Consult a radiologist if needed.",
        ],
        "differential_diagnosis": [],
        "urgency": "low",
        "follow_up": "As clinically indicated.",
    },
    "Pleural Effusion": {
        "title": "Suspected Pleural Effusion",
        "findings": "There are signs that may be consistent with fluid accumulation in the pleural space.",
        "impression": "The automated result suggests possible pleural effusion.",
        "recommendations": [
            "Consult a pulmonologist or internist.",
            "Correlate with clinical examination.",
            "Consider additional imaging as per physician's decision.",
        ],
        "differential_diagnosis": [
            "Heart failure",
            "Pneumonia",
            "Other pleural diseases",
        ],
        "urgency": "medium",
        "follow_up": "Based on symptoms and clinical evaluation.",
    },
    "Pneumonia": {
        "title": "Suspected Pneumonia",
        "findings": "Radiographic signs may be consistent with pneumonia.",
        "impression": "The automated result suggests possible pneumonia.",
        "recommendations": [
            "Consult a physician.",
            "Correlate imaging with symptoms and laboratory tests.",
            "Assess oxygen saturation if needed.",
        ],
        "differential_diagnosis": [
            "Viral pneumonia",
            "Bacterial pneumonia",
            "Pulmonary edema",
        ],
        "urgency": "high",
        "follow_up": "Based on symptom severity and physician assessment.",
    },
    "Tuberculosis": {
        "title": "Suspected Pulmonary Tuberculosis",
        "findings": "Radiographic signs may be consistent with pulmonary tuberculosis.",
        "impression": "The automated result suggests possible pulmonary tuberculosis.",
        "recommendations": [
            "Consult a pulmonologist or infectious disease specialist.",
            "Perform appropriate tests to confirm the diagnosis.",
            "Do not rely solely on the model result.",
        ],
        "differential_diagnosis": [
            "Pneumonia",
            "Other pulmonary diseases",
            "Lung tumors",
        ],
        "urgency": "high",
        "follow_up": "Urgent follow-up with a specialist.",
    },
}


class XRayModel:
    def __init__(self, model_path: str):
        logger.info("Loading EfficientNet-B4...")
        self.model = models.efficientnet_b4(weights=None)

        num_features = self.model.classifier[1].in_features
        self.model.classifier = nn.Sequential(
            nn.Dropout(0.4),
            nn.Linear(num_features, len(CLASSES)),
        )

        logger.info("Loading checkpoint: %s", model_path)
        if not os.path.exists(model_path):
            raise FileNotFoundError(
                f"Model checkpoint not found: {model_path}"
            )

        try:
            state_dict = torch.load(
                model_path,
                map_location="cpu",
                weights_only=True,
            )
        except TypeError:
            state_dict = torch.load(
                model_path,
                map_location="cpu",
            )

        self.model.load_state_dict(state_dict, strict=True)
        del state_dict

        self.model.eval()

        for parameter in self.model.parameters():
            parameter.requires_grad_(False)

        self.transform = transforms.Compose([
            transforms.Resize((IMG_SIZE, IMG_SIZE)),
            transforms.ToTensor(),
            transforms.Normalize(
                [0.485, 0.456, 0.406],
                [0.229, 0.224, 0.225],
            ),
        ])

        logger.info(
            "Model loaded. Parameters: %s",
            sum(p.numel() for p in self.model.parameters()),
        )

    def preprocess(self, image: Image.Image, device: torch.device):
        if image.mode != "RGB":
            image = image.convert("RGB")

        tensor = self.transform(image).unsqueeze(0)
        return tensor.to(device)


class GradCAM:
    def __init__(self, model: nn.Module):
        self.model = model
        self.feature_maps = None
        self.gradients = None

        self.target_layer = self.find_target_layer()
        self.target_layer.register_forward_hook(self.forward_hook)
        self.target_layer.register_full_backward_hook(self.backward_hook)

    def find_target_layer(self):
        for module in reversed(list(self.model.features.modules())):
            if isinstance(module, nn.Conv2d):
                return module
        raise RuntimeError("No Conv2d layer found.")

    def get_layer_name(self):
        for name, module in self.model.named_modules():
            if module is self.target_layer:
                return name
        return "unknown"

    def forward_hook(self, module, inputs, output):
        self.feature_maps = output

    def backward_hook(self, module, grad_input, grad_output):
        if grad_output and grad_output[0] is not None:
            self.gradients = grad_output[0]

    def clear(self):
        self.feature_maps = None
        self.gradients = None

    def generate(self, image_tensor, class_idx):
        self.clear()
        self.model.zero_grad(set_to_none=True)

        if not image_tensor.requires_grad:
            image_tensor.requires_grad_(True)

        output = self.model(image_tensor)
        score = output[0, class_idx]
        score.backward()

        if self.gradients is None or self.feature_maps is None:
            raise RuntimeError(
                "Grad-CAM hooks did not capture activations."
            )

        gradients = self.gradients[0]
        feature_maps = self.feature_maps[0]

        weights = gradients.mean(dim=(1, 2))
        cam = (weights[:, None, None] * feature_maps).sum(dim=0)
        cam = torch.relu(cam)

        cam = cam.detach().float().cpu().numpy()

        max_value = float(cam.max()) if cam.size else 0.0
        if max_value > 0:
            cam = cam / max_value

        return cam.astype(np.float32)


logger.info("Initializing inference engine...")
xray_model = XRayModel(MODEL_PATH)
grad_cam = GradCAM(xray_model.model)

# Grad-CAM stores request-specific activations in shared hooks.
# Serialize inference to prevent cross-request activation mixing.
INFERENCE_LOCK = threading.Lock()


def validate_image(file_bytes: bytes, filename: str) -> Image.Image:
    if len(file_bytes) > MAX_FILE_SIZE:
        raise ValueError("حجم الصورة يتجاوز 15MB.")

    extension = ""
    if "." in filename:
        extension = filename.rsplit(".", 1)[1].lower()

    allowed = {"jpg", "jpeg", "png", "bmp", "webp"}
    if extension not in allowed:
        raise ValueError("امتداد الصورة غير مدعوم.")

    try:
        image = Image.open(io.BytesIO(file_bytes))
        image.verify()

        image = Image.open(io.BytesIO(file_bytes))
    except UnidentifiedImageError:
        raise ValueError("الملف ليس صورة صالحة.")
    except Exception as exc:
        raise ValueError(f"خطأ في قراءة الصورة: {exc}")

    width, height = image.size

    if width < MIN_WIDTH or height < MIN_HEIGHT:
        raise ValueError(
            f"دقة الصورة منخفضة جدًا. الحد الأدنى {MIN_WIDTH}x{MIN_HEIGHT}."
        )

    if width > MAX_WIDTH or height > MAX_HEIGHT:
        raise ValueError(
            f"دقة الصورة مرتفعة جدًا. الحد الأقصى {MAX_WIDTH}x{MAX_HEIGHT}."
        )

    image = ImageOps.exif_transpose(image)
    return image.convert("RGB")


def compute_quality_metrics(image: Image.Image):
    gray = np.asarray(image.convert("L"))

    brightness = float(gray.mean())
    contrast = float(gray.std())

    sharpness = float(
        cv2.Laplacian(gray, cv2.CV_64F).var()
    )

    black_ratio = float((gray < 10).mean())
    white_ratio = float((gray > 245).mean())

    quality_score = 1.0
    issues = []

    if brightness < 20 or brightness > 220:
        issues.append("سطوع غير طبيعي")
        quality_score *= 0.7

    if contrast < 15:
        issues.append("تباين منخفض جدًا")
        quality_score *= 0.6

    if sharpness < 30:
        issues.append("الصورة قد تكون غير حادة")
        quality_score *= 0.7

    if black_ratio > 0.95 or white_ratio > 0.95:
        issues.append("مناطق فارغة كبيرة جدًا")
        quality_score *= 0.5

    return {
        "brightness": round(brightness, 2),
        "contrast": round(contrast, 2),
        "sharpness": round(sharpness, 2),
        "black_pixel_ratio": round(black_ratio, 4),
        "white_pixel_ratio": round(white_ratio, 4),
        "quality_score": round(quality_score, 2),
        "issues": issues,
        "is_acceptable": quality_score >= 0.5,
    }


def check_confidence(probs, threshold=0.60):
    sorted_probs, _ = torch.sort(probs, descending=True)

    top1 = float(sorted_probs[0].item())
    top2 = float(sorted_probs[1].item())
    margin = top1 - top2

    flags = []

    if top1 < 0.40:
        level = "very_low"
        flags.append("very_low_confidence")
    elif top1 < 0.60:
        level = "low"
        flags.append("low_confidence")
    elif top1 < 0.85:
        level = "medium"
    else:
        level = "high"

    if margin < 0.10:
        flags.append("ambiguous_prediction")
        level = "ambiguous"

    reliable = top1 >= threshold and margin >= 0.10

    return {
        "confidence_level": level,
        "flags": flags,
        "top1_probability": top1,
        "top2_probability": top2,
        "margin": margin,
        "is_reliable": reliable,
    }


def _encode_visualization_bytes(image_array, fmt="png", jpeg_quality=88):
    """Encode a visualization to binary bytes for Storage upload only.
    This is intentionally bytes, never Base64. The bytes stay server-side and
    are uploaded directly to private Cloudflare R2 by database.save_analysis().
    """
    array = np.asarray(image_array)
    if array.dtype != np.uint8:
        if array.size and float(array.max()) <= 1.0:
            array = (array * 255).clip(0, 255).astype(np.uint8)
        else:
            array = array.clip(0, 255).astype(np.uint8)

    fmt = (fmt or "png").lower().lstrip(".")
    if fmt == "png":
        encode_image = array if array.ndim == 2 else cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
        ext = ".png"
        params = [cv2.IMWRITE_PNG_COMPRESSION, 3]
    else:
        encode_image = array if array.ndim == 2 else cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
        ext = ".jpg"
        params = [cv2.IMWRITE_JPEG_QUALITY, int(jpeg_quality)]

    ok, encoded = cv2.imencode(ext, encode_image, params)
    if not ok:
        raise RuntimeError("Failed to encode visualization.")
    return encoded.tobytes()


def enhance_heatmap(
    heatmap,
    percentile_low=5,
    percentile_high=95,
    gamma=0.8,
):
    heatmap = np.asarray(heatmap, dtype=np.float32)

    low = np.percentile(heatmap, percentile_low)
    high = np.percentile(heatmap, percentile_high)

    enhanced = heatmap.copy()
    enhanced[enhanced < low] = 0
    enhanced = np.clip(enhanced, 0, high)

    minimum = float(enhanced.min())
    maximum = float(enhanced.max())

    if maximum > minimum:
        enhanced = (enhanced - minimum) / (maximum - minimum)

    enhanced = np.power(enhanced, gamma)
    return enhanced.astype(np.float32)


def create_overlay(image_rgb, heatmap, alpha=0.6):
    if heatmap.shape != image_rgb.shape[:2]:
        heatmap = cv2.resize(
            heatmap,
            (image_rgb.shape[1], image_rgb.shape[0]),
            interpolation=cv2.INTER_LINEAR,
        )

    colored = cv2.applyColorMap(
        (heatmap * 255).astype(np.uint8),
        cv2.COLORMAP_INFERNO,
    )

    colored = cv2.cvtColor(colored, cv2.COLOR_BGR2RGB).astype(
        np.float32
    ) / 255.0

    overlay = colored * alpha + image_rgb * (1 - alpha)
    return np.clip(overlay, 0, 1)


def generate_visualizations(image, input_tensor, pred_class):
    original_np = np.asarray(image).astype(np.float32) / 255.0

    with torch.enable_grad():
        try:
            heatmap_raw = grad_cam.generate(input_tensor, pred_class)
        except Exception as exc:
            logger.warning("Grad-CAM failed: %s", exc)
            heatmap_raw = np.zeros(
                (IMG_SIZE, IMG_SIZE),
                dtype=np.float32,
            )

    active_ratio = float((heatmap_raw > 0.1).mean())
    mean_activation = float(heatmap_raw.mean())
    std_activation = float(heatmap_raw.std())

    heatmap = enhance_heatmap(
        heatmap_raw,
        percentile_low=5,
        percentile_high=95,
        gamma=0.8,
    )

    overlay = create_overlay(
        original_np,
        heatmap,
        alpha=0.6,
    )

    heatmap_resized_full = cv2.resize(
        heatmap,
        (original_np.shape[1], original_np.shape[0]),
        interpolation=cv2.INTER_LINEAR,
    )

    heatmap_color = cv2.applyColorMap(
        (heatmap_resized_full * 255).astype(np.uint8),
        cv2.COLORMAP_INFERNO,
    )
    heatmap_color = cv2.cvtColor(
        heatmap_color,
        cv2.COLOR_BGR2RGB,
    )

    # Keep the original contour behavior: threshold the 380x380 CAM,
    # then draw contours on the original-resolution image.
    heatmap_resized = cv2.resize(
        heatmap,
        (IMG_SIZE, IMG_SIZE),
        interpolation=cv2.INTER_LINEAR,
    )

    _, threshold = cv2.threshold(
        (heatmap_resized * 255).astype(np.uint8),
        50,
        255,
        cv2.THRESH_BINARY,
    )

    contours, _ = cv2.findContours(
        threshold,
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE,
    )

    # Scale contours to the original image dimensions.
    scale_x = original_np.shape[1] / IMG_SIZE
    scale_y = original_np.shape[0] / IMG_SIZE

    contour_image = (original_np * 255).astype(np.uint8)

    if contours:
        scaled_contours = []
        for contour in contours:
            c = contour.astype(np.float32)
            c[:, :, 0] *= scale_x
            c[:, :, 1] *= scale_y
            scaled_contours.append(np.round(c).astype(np.int32))

        cv2.drawContours(
            contour_image,
            scaled_contours,
            -1,
            (0, 255, 255),
            3,
        )

    contour_image = contour_image.astype(np.float32) / 255.0

    threshold_overlay = original_np.copy()

    mask = cv2.resize(
        (heatmap > 0.5).astype(np.float32),
        (original_np.shape[1], original_np.shape[0]),
        interpolation=cv2.INTER_NEAREST,
    )

    threshold_overlay[mask > 0.5] = (1.0, 0.0, 0.0)

    comparison = np.hstack([original_np, overlay])

    images = {
        # Binary bytes are used only for server-side Storage upload.
        "heatmap_enhanced": _encode_visualization_bytes(heatmap_color, "png"),
        "overlay": _encode_visualization_bytes(overlay, "jpg"),
        "contour": _encode_visualization_bytes(contour_image, "jpg"),
        "threshold": _encode_visualization_bytes(threshold_overlay, "jpg"),
        "comparison": _encode_visualization_bytes(comparison, "jpg"),
    }

    statistics = {
        "active_ratio": round(active_ratio, 6),
        "mean_activation": round(mean_activation, 6),
        "std_activation": round(std_activation, 6),
    }

    grad_cam.clear()
    return images, statistics


def generate_medical_report(
    pred_class,
    confidence,
    probs,
    quality,
    confidence_check,
    language="ar",
):
    templates = REPORT_TEMPLATES_EN if language == "en" else REPORT_TEMPLATES_AR

    class_name = CLASSES[pred_class]

    if class_name not in templates:
        templates = REPORT_TEMPLATES_AR

    template = templates[class_name]

    sorted_probs, sorted_indices = torch.sort(
        probs,
        descending=True,
    )

    top3 = []
    for i in range(min(3, len(sorted_probs))):
        idx = int(sorted_indices[i].item())
        probability = float(sorted_probs[i].item())

        top3.append({
            "disease": CLASSES[idx],
            "probability": round(probability, 6),
            "percentage": f"{probability * 100:.2f}%",
        })

    now = datetime.utcnow()
    report_id = "XR-" + now.strftime("%Y%m%d%H%M%S%f")[:-3]

    return {
        "report_id": report_id,
        "generated_at": now.isoformat() + "Z",
        "language": language,
        "title": template["title"],
        "prediction": {
            "primary_diagnosis": class_name,
            "confidence": round(confidence, 6),
            "confidence_percentage": f"{confidence * 100:.2f}%",
            "confidence_level": confidence_check["confidence_level"],
        },
        "findings": template["findings"],
        "impression": template["impression"],
        "recommendations": template["recommendations"],
        "differential_diagnosis": template["differential_diagnosis"],
        "urgency": template["urgency"],
        "follow_up": template["follow_up"],
        "quality_metrics": quality,
        "confidence_flags": confidence_check["flags"],
        "reliability": {
            "is_reliable": confidence_check["is_reliable"],
            "reasons": (
                confidence_check["flags"]
                if confidence_check["flags"]
                else (
                    ["No clear issues"]
                    if language == "en"
                    else ["لا توجد مشاكل واضحة"]
                )
            ),
        },
        "top_3_predictions": top3,
        "disclaimer": (
            "This report was automatically generated by an AI system for assistance only, "
            "and does not replace evaluation by a physician or radiologist."
            if language == "en"
            else
            "هذا التقرير تم إنشاؤه آليًا بواسطة نظام ذكاء اصطناعي للمساعدة فقط، "
            "ولا يغني عن تقييم الطبيب أو أخصائي الأشعة."
        ),
    }



def generate_bilingual_reports(pred_class, confidence, probs, quality, confidence_check):
    report_ar = generate_medical_report(pred_class, confidence, probs, quality, confidence_check, language="ar")
    report_en = generate_medical_report(pred_class, confidence, probs, quality, confidence_check, language="en")
    report_en["report_id"] = report_ar["report_id"]
    report_ar["language"] = "ar"
    report_en["language"] = "en"
    return report_ar, report_en

def predict_bytes(file_bytes: bytes, filename="xray.jpg", language="ar"):
    """Run EfficientNet-B4 + Grad-CAM inference and build bilingual reports."""
    request_tag = f"{filename}|{len(file_bytes)}bytes"

    if language not in {"ar", "en"}:
        language = "ar"

    logger.info("INFERENCE START %s language=%s", request_tag, language)

    # ------------------------------------------------------------
    # Image validation / quality
    # ------------------------------------------------------------
    logger.info("INFERENCE STEP image validation START %s", request_tag)
    image = validate_image(file_bytes, filename)
    logger.info(
        "INFERENCE STEP image validation OK %s size=%sx%s mode=%s",
        request_tag,
        image.width,
        image.height,
        image.mode,
    )

    quality = compute_quality_metrics(image)
    logger.info(
        "INFERENCE STEP quality OK %s score=%s acceptable=%s",
        request_tag,
        quality.get("quality_score"),
        quality.get("is_acceptable"),
    )

    if not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA/ZeroGPU is not available. Create this Space with ZeroGPU hardware."
        )

    gpu = torch.device("cuda")
    input_tensor = None
    outputs = None
    probabilities = None

    # Grad-CAM uses shared hooks; serialize requests.
    with INFERENCE_LOCK:
        try:
            # --------------------------------------------------------
            # Model / prediction
            # --------------------------------------------------------
            logger.info("INFERENCE STEP model-to-GPU START %s", request_tag)
            xray_model.model.to(gpu)
            xray_model.model.eval()

            input_tensor = xray_model.preprocess(image, gpu)

            start_time = datetime.utcnow()

            logger.info("INFERENCE STEP forward START %s", request_tag)
            with torch.inference_mode():
                outputs = xray_model.model(input_tensor)
                probabilities = F.softmax(outputs, dim=1)[0]
                confidence_tensor, pred_tensor = torch.max(
                    probabilities,
                    dim=0,
                )

            pred_class = int(pred_tensor.item())
            confidence = float(confidence_tensor.item())

            logger.info(
                "INFERENCE STEP forward OK %s class=%s confidence=%.6f",
                request_tag,
                CLASSES[pred_class],
                confidence,
            )

            confidence_check = check_confidence(probabilities)

            # --------------------------------------------------------
            # Grad-CAM / visualizations
            # --------------------------------------------------------
            logger.info("INFERENCE STEP Grad-CAM START %s", request_tag)
            try:
                visualizations, heatmap_stats = generate_visualizations(
                    image,
                    input_tensor,
                    pred_class,
                )
                logger.info("INFERENCE STEP Grad-CAM OK %s", request_tag)
            except Exception as exc:
                logger.exception(
                    "INFERENCE STEP Grad-CAM FAILED %s: %s",
                    request_tag,
                    exc,
                )
                visualizations = {
                    "heatmap_enhanced": "",
                    "overlay": "",
                    "contour": "",
                    "threshold": "",
                    "comparison": "",
                }
                heatmap_stats = {
                    "active_ratio": 0.0,
                    "mean_activation": 0.0,
                    "std_activation": 0.0,
                }

            # --------------------------------------------------------
            # Janus-Pro-CXR report + deterministic parser
            # --------------------------------------------------------
            logger.info("INFERENCE STEP Janus reports START %s", request_tag)
            janus_payload = {"enabled": False, "report_ar": "", "report_en": "", "requested_raw_text": ""}
            try:
                janus_payload = janus_generate_reports(image, requested_language=language)
            except Exception as janus_exc:
                logger.exception("JANUS REPORT FAILED %s: %s", request_tag, janus_exc)

            template_ar, template_en = generate_bilingual_reports(
                pred_class, confidence, probabilities, quality, confidence_check
            )

            def merge_janus(template, raw_text, lang):
                parsed = parse_radiology_text(raw_text) if raw_text else {}
                if not raw_text:
                    template["report_source"] = "EFFICIENTNET_TEMPLATE_FALLBACK"
                    template["janus_raw_text"] = ""
                    return template
                if parsed.get("findings"):
                    template["findings"] = parsed["findings"]
                if parsed.get("impression"):
                    template["impression"] = parsed["impression"]
                if parsed.get("recommendations"):
                    template["recommendations"] = parsed["recommendations"]
                template["report_source"] = "JANUS_PRO_CXR"
                template["report_model"] = "ZrH42/Janus-Pro-CXR-Final"
                template["janus_raw_text"] = raw_text
                template["janus_parsed"] = parsed
                return template

            arabic_raw = janus_payload.get("report_ar", "") or ""
            english_raw = janus_payload.get("report_en", "") or ""
            arabic_verified = bool(janus_payload.get("arabic_verified"))
            english_verified = bool(janus_payload.get("english_verified"))

            # Never allow an unverified language payload into the corresponding report.
            if not arabic_verified or not re.search(r"[\u0600-\u06FF]", arabic_raw):
                if arabic_raw:
                    logger.warning("JANUS Arabic output failed language validation; using Arabic template fallback")
                arabic_raw = ""
                arabic_verified = False

            if not english_verified or not re.search(r"[A-Za-z]", english_raw):
                if english_raw:
                    logger.warning("JANUS English output failed language validation; using English template fallback")
                english_raw = ""
                english_verified = False

            # Merge verified native Janus outputs into their matching language.
            # Never cross-copy English into Arabic.
            report_ar = merge_janus(template_ar, arabic_raw, "ar")
            report_en = merge_janus(template_en, english_raw, "en")

            # Arabic strategy:
            # 1) Prefer a genuinely Arabic Janus report if one exists.
            # 2) Otherwise translate the VERIFIED Janus English structured report.
            # 3) Accept the translation only after conservative script/number checks.
            # 4) Otherwise use the existing Arabic EfficientNet template.
            arabic_translation = {
                "enabled": False,
                "verified": False,
                "model": None,
                "findings": "",
                "impression": "",
                "recommendations": [],
            }

            if not arabic_verified and english_verified:
                try:
                    english_parsed_for_translation = parse_radiology_text(english_raw)
                    arabic_translation = translate_report_en_to_ar(english_parsed_for_translation)
                    if arabic_translation.get("verified"):
                        report_ar["findings"] = arabic_translation.get("findings", "") or report_ar.get("findings", "")
                        report_ar["impression"] = arabic_translation.get("impression", "") or report_ar.get("impression", "")
                        if arabic_translation.get("recommendations"):
                            report_ar["recommendations"] = arabic_translation["recommendations"]
                        report_ar["report_source"] = "JANUS_EN_TRANSLATED_TO_ARABIC"
                        report_ar["report_model"] = arabic_translation.get("model") or "Helsinki-NLP/opus-mt-en-ar"
                        report_ar["janus_raw_text"] = ""
                        report_ar["translation_source"] = "JANUS_PRO_CXR_ENGLISH"
                        report_ar["translation_verified"] = True
                    else:
                        logger.warning(
                            "Arabic controlled translation rejected: script_valid=%s numbers_preserved=%s gibberish_free=%s terms=%s",
                            arabic_translation.get("script_valid"),
                            arabic_translation.get("numbers_preserved"),
                            arabic_translation.get("gibberish_free"),
                            arabic_translation.get("required_terms", {}).get("all_present"),
                        )
                except Exception as translation_exc:
                    logger.exception("Arabic translation failed: %s", translation_exc)
                    arabic_translation = {
                        "enabled": True,
                        "verified": False,
                        "model": "Helsinki-NLP/opus-mt-en-ar",
                        "error": str(translation_exc),
                    }

            # If Janus English was verified but Arabic conversion failed, do not
            # silently present an unrelated EfficientNet Arabic report as if it
            # were a translation of the same analysis. Mark the result for review.
            arabic_translation_failed = bool(
                (language == "ar") and english_verified and not report_ar.get("translation_verified")
            )
            if arabic_translation_failed:
                # Do NOT expose an unrelated EfficientNet Arabic template when a
                # verified Janus English report exists. Keep the Janus source
                # traceable and mark the Arabic presentation as review-required.
                report_ar["findings"] = ""
                report_ar["impression"] = ""
                report_ar["recommendations"] = []
                report_ar["report_source"] = "ARABIC_TRANSLATION_REQUIRES_REVIEW"
                report_ar["translation_verified"] = False
                report_ar["translation_source"] = "JANUS_PRO_CXR_ENGLISH"
                report_ar["janus_source_report_en"] = english_raw

            report = report_en if language == "en" else report_ar
            logger.info("INFERENCE STEP reports OK %s source_ar=%s source_en=%s", request_tag, report_ar.get("report_source"), report_en.get("report_source"))

            all_predictions = []
            for index, probability in enumerate(
                probabilities.detach().float().cpu().tolist()
            ):
                all_predictions.append({
                    "class": CLASSES[index],
                    "probability": round(probability, 6),
                    "percentage": round(probability * 100, 2),
                })

            elapsed_ms = (
                datetime.utcnow() - start_time
            ).total_seconds() * 1000.0

            report_integrity = {
                "requested_language": language,
                "arabic": {
                    "language": "ar",
                    # This is the verification state of the FINAL Arabic report,
                    # not merely native Janus Arabic generation.
                    "verified": bool(
                        arabic_verified or report_ar.get("translation_verified") is True
                    ),
                    "script_valid": bool(re.search(r"[\u0600-\u06FF]", report_ar.get("findings", "") + " " + report_ar.get("impression", ""))),
                    "source": report_ar.get("report_source", "EFFICIENTNET_TEMPLATE_FALLBACK"),
                    "native_janus_verified": bool(arabic_verified),
                    "translation_verified": bool(report_ar.get("translation_verified")),
                    "requires_manual_review": bool(arabic_translation_failed),
                },
                "english": {
                    "language": "en",
                    "verified": bool(english_verified),
                    "script_valid": bool(re.search(r"[A-Za-z]", report_en.get("findings", "") + " " + report_en.get("impression", ""))),
                    "source": report_en.get("report_source", "EFFICIENTNET_TEMPLATE_FALLBACK"),
                },
                "cross_language_fallback": False,
                "arabic_translation": {
                    "attempted": bool(not arabic_verified and english_verified),
                    "verified": bool(arabic_translation.get("verified")),
                    "model": arabic_translation.get("model"),
                    "script_valid": bool(arabic_translation.get("script_valid", False)),
                    "numbers_preserved": bool(arabic_translation.get("numbers_preserved", False)),
                    "item_count_preserved": bool(arabic_translation.get("item_count_preserved", False)),
                    "negation_preserved": bool(arabic_translation.get("negation_preserved", False)),
                    "required_terms": arabic_translation.get("required_terms", {}),
                    "gibberish_free": bool(arabic_translation.get("gibberish_free", False)),
                    "translation_methods": arabic_translation.get("translation_methods", []),
                    "translated_text_preview": arabic_translation.get("translated_text_preview", ""),
                    "warning": arabic_translation.get("warning"),
                },
                "consistency": {
                    "english_source_available": bool(english_verified),
                    "english_source": "JANUS_PRO_CXR" if english_verified else "UNAVAILABLE",
                    "same_analysis_source": bool(english_verified and (arabic_translation.get("verified") or arabic_verified)),
                    "arabic_source": (
                        "JANUS_PRO_CXR_AR" if arabic_verified
                        else "JANUS_EN_TRANSLATED_TO_ARABIC" if arabic_translation.get("verified")
                        else "ARABIC_TRANSLATION_REQUIRES_REVIEW" if english_verified
                        else "EFFICIENTNET_TEMPLATE_FALLBACK"
                    ),
                    "translation_verified": bool(arabic_translation.get("verified")),
                    "requires_manual_review": bool(arabic_translation_failed),
                },
            }

            # Independent conservative checks across the final Arabic/English reports.
            # These checks flag review risks; they do not establish clinical correctness.
            bilingual_checks = check_bilingual_reports(
                report_ar=report_ar,
                report_en=report_en,
                classifier=CLASSES[pred_class],
                confidence=confidence,
                quality=quality,
                translation_meta=report_integrity.get("arabic_translation", {}),
            )
            report_integrity["bilingual_quality"] = bilingual_checks
            report_integrity["manual_review_required"] = bool(bilingual_checks["manual_review_required"])
            report_integrity["review_status"] = bilingual_checks["status"]

            response = {
                "status": "success",
                "language": language,
                "analysis_id": report_ar["report_id"],
                "medical_report": report,
                "medical_report_ar": report_ar,
                "medical_report_en": report_en,
                "report_languages": ["ar", "en"],
                "report_integrity": report_integrity,
                "predictions": all_predictions,
                "visualizations": visualizations,
                "heatmap_stats": heatmap_stats,
                "quality_metrics": quality,
                "model_info": {
                    "architecture": "EfficientNet-B4",
                    "input_size": IMG_SIZE,
                    "classes": len(CLASSES),
                    "device": "cuda",
                    "version": "2.0.0-kaggle-flask",
                    "inference_time_ms": round(elapsed_ms, 2),
                    "report_model": "ZrH42/Janus-Pro-CXR-Final" if janus_payload.get("enabled") else "template_fallback",
                    "report_source": report.get("report_source", "EFFICIENTNET_TEMPLATE_FALLBACK"),
                },
                "janus": {
                    "enabled": bool(janus_payload.get("enabled")),
                    "model": "ZrH42/Janus-Pro-CXR-Final" if janus_payload.get("enabled") else None,
                    "raw_text": janus_payload.get("requested_raw_text", ""),
                    "report_ar_raw": janus_payload.get("report_ar", ""),
                    "report_en_raw": janus_payload.get("report_en", ""),
                    "arabic_verified": arabic_verified,
                    "english_verified": english_verified,
                    "arabic_translation": arabic_translation,
                },
            }

            logger.info(
                "INFERENCE SUCCESS %s class=%s confidence=%.6f time_ms=%.2f",
                request_tag,
                CLASSES[pred_class],
                confidence,
                elapsed_ms,
            )

            return response

        except Exception as exc:
            logger.exception(
                "INFERENCE FAILED %s type=%s message=%s",
                request_tag,
                type(exc).__name__,
                exc,
            )
            raise

        finally:
            # Always release request-specific tensors and return weights to CPU.
            try:
                if input_tensor is not None:
                    del input_tensor
                if outputs is not None:
                    del outputs
                if probabilities is not None:
                    del probabilities
            except Exception:
                pass

            try:
                xray_model.model.to("cpu")
            except Exception:
                logger.exception("INFERENCE CLEANUP: failed moving model to CPU")

            try:
                grad_cam.clear()
            except Exception:
                pass

            gc.collect()

            try:
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except Exception:
                pass

            logger.info("INFERENCE CLEANUP COMPLETE %s", request_tag)


def health_info():
    return {
        "status": "healthy",
        "service": "Clinical AI X-Ray API",
        "model": "EfficientNet-B4",
        "device": "cuda/ZeroGPU",
        "classes": len(CLASSES),
        "gradcam": True,
        "version": "1.0.0-zerogpu",
    }


def classes_info():
    return {
        "classes": CLASSES,
        "count": len(CLASSES),
    }


def root_info():
    return {
        "name": "Clinical AI X-Ray API",
        "version": "1.0.0-zerogpu",
        "status": "online",
        "endpoints": {
            "health": "Gradio/API health endpoint",
            "classes": "Gradio/API classes endpoint",
            "predict": "Gradio/API predict endpoint",
        },
    }