app.py

import streamlit as st
from transformers import pipeline
from PIL import Image, ImageDraw
import torch

st.set_page_config(
    page_title="Fraktur Detektion",
    layout="wide",
    initial_sidebar_state="collapsed"
)

st.markdown("""
<style>
    /* Reset et base */
    .stApp {
        background-color: var(--background-color) !important;
        padding: 0 !important;
        overflow: hidden !important;
        height: 100vh !important;
    }

    /* Variables de thème */
    [data-theme="light"] {
        --background-color: #ffffff;
        --text-color: #1f2937;
        --border-color: #e5e7eb;
        --secondary-bg: #f3f4f6;
    }

    [data-theme="dark"] {
        --background-color: #1f2937;
        --text-color: #f3f4f6;
        --border-color: #4b5563;
        --secondary-bg: #374151;
    }

    /* Layout principal */
    .block-container {
        padding: 0.5rem !important;
        max-width: 100% !important;
    }

    /* Contrôles et upload */
    .uploadedFile {
        border: 1px dashed var(--border-color);
        border-radius: 0.375rem;
        padding: 0.25rem;
        background: var(--secondary-bg);
    }

    /* Ajustement des colonnes */
    [data-testid="column"] {
        padding: 0 0.5rem !important;
    }

    /* Images adaptatives */
    .stImage > img {
        width: 100% !important;
        height: auto !important;
        max-height: 400px !important;
        object-fit: contain !important;
    }

    /* Résultats */
    .result-box {
        padding: 0.375rem;
        border-radius: 0.375rem;
        margin: 0.25rem 0;
        background: var(--secondary-bg);
        border: 1px solid var(--border-color);
        color: var(--text-color);
    }

    /* Titres */
    h2, h3 {
        margin: 0 !important;
        padding: 0.5rem 0 !important;
        font-size: 1rem !important;
        color: var(--text-color) !important;
    }

    /* Nettoyage des éléments inutiles */
    #MainMenu, footer, header, .viewerBadge_container__1QSob, .stDeployButton {
        display: none !important;
    }

    /* Ajustements espacement */
    div[data-testid="stVerticalBlock"] {
        gap: 0.5rem !important;
    }

    .element-container {
        margin: 0.25rem 0 !important;
    }
</style>
""", unsafe_allow_html=True)

@st.cache_resource
def load_models():
    return {
        "KnochenAuge": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
        "KnochenWächter": pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray"),
        "RöntgenMeister": pipeline("image-classification", 
            model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388")
    }

def translate_label(label):
    translations = {
        "fracture": "Knochenbruch",
        "no fracture": "Kein Bruch",
        "normal": "Normal",
        "abnormal": "Auffällig"
    }
    return translations.get(label.lower(), label)

def draw_boxes(image, predictions):
    draw = ImageDraw.Draw(image)
    for pred in predictions:
        box = pred['box']
        label = f"{translate_label(pred['label'])} ({pred['score']:.2%})"
        color = "#2563eb" if pred['score'] > 0.7 else "#eab308"

        draw.rectangle(
            [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
            outline=color,
            width=2
        )

        # Label plus compact
        text_bbox = draw.textbbox((box['xmin'], box['ymin']-15), label)
        draw.rectangle(text_bbox, fill=color)
        draw.text((box['xmin'], box['ymin']-15), label, fill="white")
    return image

def main():
    models = load_models()

    # Disposition en deux colonnes principales
    col1, col2 = st.columns([1, 2])

    with col1:
        st.markdown("### 📤 Röntgenbild Upload")
        uploaded_file = st.file_uploader("", type=['png', 'jpg', 'jpeg'])

        if uploaded_file:
            conf_threshold = st.slider(
                "Konfidenzschwelle",
                min_value=0.0, max_value=1.0,
                value=0.60, step=0.05
            )

    with col2:
        if uploaded_file:
            image = Image.open(uploaded_file)

            st.markdown("### 🔍 Meinung der KI-Experten")

            # Analyse avec KnochenAuge (localisation)
            st.markdown("#### 👁️ Das KnochenAuge - Lokalisation")
            predictions = models["KnochenAuge"](image)
            filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]

            if filtered_preds:
                result_image = image.copy()
                result_image = draw_boxes(result_image, filtered_preds)
                st.image(result_image, use_container_width=True)

                # Affichage des résultats supplémentaires uniquement si des fractures sont détectées
                st.markdown("#### 🎯 KI-Analyse")
                col_left, col_right = st.columns(2)

                with col_left:
                    st.markdown("**🛡️ Der KnochenWächter**")
                    predictions = models["KnochenWächter"](image)
                    for pred in predictions:
                        if pred['score'] >= conf_threshold:
                            score_color = "#22c55e" if pred['score'] > 0.7 else "#eab308"
                            st.markdown(f"""
                                <div class='result-box'>
                                    <span style='color: {score_color}; font-weight: 500;'>
                                        {pred['score']:.1%}
                                    </span> - {translate_label(pred['label'])}
                                </div>
                            """, unsafe_allow_html=True)

                with col_right:
                    st.markdown("**🎓 Der RöntgenMeister**")
                    predictions = models["RöntgenMeister"](image)
                    for pred in predictions:
                        if pred['score'] >= conf_threshold:
                            score_color = "#22c55e" if pred['score'] > 0.7 else "#eab308"
                            st.markdown(f"""
                                <div class='result-box'>
                                    <span style='color: {score_color}; font-weight: 500;'>
                                        {pred['score']:.1%}
                                    </span> - {translate_label(pred['label'])}
                                </div>
                            """, unsafe_allow_html=True)
            else:
                st.info("Kein Bruch erkannt.")
        else:
            st.info("Bitte laden Sie ein Röntgenbild hoch (JPEG, PNG)")

    # Script pour la synchronisation du thème
    st.markdown("""
        <script>
            function updateTheme(isDark) {
                document.documentElement.setAttribute('data-theme', isDark ? 'dark' : 'light');
            }

            window.addEventListener('message', function(e) {
                if (e.data.type === 'theme-change') {
                    updateTheme(e.data.theme === 'dark');
                }
            });

            updateTheme(window.matchMedia('(prefers-color-scheme: dark)').matches);
        </script>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()