Spaces:

yassonee
/

chest

Running

File size: 9,296 Bytes

from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse
from transformers import pipeline
from PIL import Image
import io
import uvicorn
import base64

app = FastAPI()

# Chargement des modèles
def load_models():
    return {
        "chest_classifier": pipeline("image-classification", model="codewithdark/vit-chest-xray")
    }

models = load_models()

# Liste des maladies dans l'ordre des labels du modèle
LABEL_MAP = {
    'LABEL_0': 'Kardiomegalie',  # Élargissement du cœur
    'LABEL_1': 'Ödem',           # Œdème
    'LABEL_2': 'Konsolidierung', # Consolidation
    'LABEL_3': 'Lungenentzündung', # Pneumonie
    'LABEL_4': 'Kein Befund'     # Aucune anomalie
}

def translate_label(label):
    # Si le label est au format "LABEL_X", on le traduit
    if isinstance(label, str) and label in LABEL_MAP:
        return LABEL_MAP[label]
    # Sinon on retourne le label original avec un avertissement
    return f"Unbekannt ({label})"
    return translations.get(label, label)

def image_to_base64(image):
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"

COMMON_STYLES = """
    body {
        font-family: system-ui, -apple-system, sans-serif;
        background: #f0f2f5;
        margin: 0;
        padding: 20px;
        color: #1a1a1a;
    }
    ::-webkit-scrollbar {
        width: 8px;
        height: 8px;
    }
    
    ::-webkit-scrollbar-track {
        background: transparent;
    }
    
    ::-webkit-scrollbar-thumb {
        background-color: rgba(156, 163, 175, 0.5);
        border-radius: 4px;
    }
    
    .container {
        max-width: 1200px;
        margin: 0 auto;
        background: white;
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    .button {
        background: #2d2d2d;
        color: white;
        border: none;
        padding: 12px 30px;
        border-radius: 8px;
        cursor: pointer;
        font-size: 1.1em;
        transition: all 0.3s ease;
        position: relative;
    }
    .button:hover {
        background: #404040;
    }
    @keyframes progress {
        0% { width: 0; }
        100% { width: 100%; }
    }
    .button-progress {
        position: absolute;
        bottom: 0;
        left: 0;
        height: 4px;
        background: rgba(255, 255, 255, 0.5);
        width: 0;
    }
    .button:active .button-progress {
        animation: progress 2s linear forwards;
    }
    img {
        max-width: 100%;
        height: auto;
        border-radius: 8px;
    }
    @keyframes blink {
        0% { opacity: 1; }
        50% { opacity: 0; }
        100% { opacity: 1; }
    }
    #loading {
        display: none;
        color: white;
        margin-top: 10px;
        animation: blink 1s infinite;
        text-align: center;
    }
"""

@app.get("/", response_class=HTMLResponse)
async def main():
    content = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>Thorax-Röntgen Analyse</title>
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <style>
            {COMMON_STYLES}
            
            .upload-section {{
                background: #2d2d2d;
                padding: 40px;
                border-radius: 12px;
                margin: 20px 0;
                text-align: center;
                border: 2px dashed #404040;
                transition: all 0.3s ease;
                color: white;
            }}
            .upload-section:hover {{
                border-color: #555;
            }}
            input[type="file"] {{
                width: 0.1px;
                height: 0.1px;
                opacity: 0;
                overflow: hidden;
                position: absolute;
                z-index: -1;
            }}
            .file-upload-label {{
                display: inline-block;
                padding: 12px 30px;
                background: #404040;
                color: white;
                border-radius: 8px;
                cursor: pointer;
                font-size: 1.1em;
                transition: all 0.3s ease;
                margin: 20px 0;
            }}
            .file-upload-label:hover {{
                background: #555;
            }}
        </style>
    </head>
    <body>
        <div class="container">
            <div class="upload-section">
                <form action="/analyze" method="post" enctype="multipart/form-data" onsubmit="document.getElementById('loading').style.display = 'block';">
                    <div>
                        <label for="file-upload" class="file-upload-label">
                            Röntgenbild auswählen
                        </label>
                        <input id="file-upload" type="file" name="file" accept="image/*" required>
                    </div>
                    <button type="submit" class="button">
                        Analysieren
                        <div class="button-progress"></div>
                    </button>
                    <div id="loading">Wird geladen...</div>
                </form>
            </div>
        </div>
    </body>
    </html>
    """
    return content

@app.post("/analyze", response_class=HTMLResponse)
async def analyze_file(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        
        predictions = models["chest_classifier"](image)
        result_image_b64 = image_to_base64(image)
        
        results_html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Ergebnisse</title>
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <style>
                {COMMON_STYLES}
                
                .results-grid {{
                    display: grid;
                    grid-template-columns: 1fr 1fr;
                    gap: 20px;
                    margin-top: 20px;
                }}
                .result-box {{
                    background: white;
                    padding: 20px;
                    border-radius: 12px;
                    margin: 10px 0;
                    border: 1px solid #e9ecef;
                }}
                .score-high {{ 
                    color: #0066cc; 
                    font-weight: bold;
                }}
                .score-medium {{ 
                    color: #ffa500; 
                    font-weight: bold;
                }}
                .back-button {{
                    display: inline-block;
                    text-decoration: none;
                    margin-top: 20px;
                }}
                h3 {{
                    color: #0066cc;
                    margin-top: 0;
                }}
                @media (max-width: 768px) {{
                    .results-grid {{
                        grid-template-columns: 1fr;
                    }}
                }}
            </style>
        </head>
        <body>
            <div class="container">
                <div class="results-grid">
                    <div class="result-box">
                        <h3>Analyse-Ergebnisse</h3>
        """
        
        for pred in predictions:
            confidence_class = "score-high" if pred['score'] > 0.7 else "score-medium"
            results_html += f"""
                <div>
                    <span class="{confidence_class}">{pred['score']:.1%}</span> - 
                    {translate_label(pred['label'])}
                </div>
            """
        
        results_html += f"""
                    </div>
                    <div class="result-box">
                        <h3>Röntgenbild</h3>
                        <img src="{result_image_b64}" alt="Analysiertes Röntgenbild">
                    </div>
                </div>
                
                <a href="/" class="button back-button">
                    ← Zurück
                    <div class="button-progress"></div>
                </a>
            </div>
        </body>
        </html>
        """
        
        return results_html
        
    except Exception as e:
        return f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Fehler</title>
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <style>
                {COMMON_STYLES}
                .error-box {{
                    background: #fee2e2;
                    border: 1px solid #ef4444;
                    padding: 20px;
                    border-radius: 8px;
                    margin: 20px 0;
                }}
            </style>
        </head>
        <body>
            <div class="container">
                <div class="error-box">
                    <h3>Fehler</h3>
                    <p>{str(e)}</p>
                </div>
                <a href="/" class="button back-button">
                    ← Zurück
                    <div class="button-progress"></div>
                </a>
            </div>
        </body>
        </html>
        """

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)