app.py

import gradio as gr
import torch
from PIL import Image
import numpy as np
from torchvision import models
from transformers import ViTForImageClassification
from torch import nn
from torch.cuda.amp import autocast
import os

# Global configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Label mapping (HAM10K)
label_mapping = {
    0: "Меланома",
    1: "Меланоцитарный невус",
    2: "Базальноклеточная карцинома",
    3: "Актинический кератоз",
    4: "Доброкачественная кератоза",
    5: "Дерматофиброма",
    6: "Сосудистые поражения"
}

# Model paths
CHECKPOINTS_PATH = os.getenv("CHECKPOINTS_PATH", "./checkpoints")

# Model definitions
def get_efficientnet():
    model = models.efficientnet_v2_s(weights="IMAGENET1K_V1")
    model.classifier[1] = nn.Linear(1280, 7)
    return model.to(device)

def get_deit():
    model = ViTForImageClassification.from_pretrained(
        'facebook/deit-base-patch16-224',
        num_labels=7,
        ignore_mismatched_sizes=True
    )
    return model.to(device)

# Transforms
def transform_image(image):
    """Transform PIL image to model input format"""
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])
    return transform(image).unsqueeze(0).to(device)

# Model Handler
class ModelHandler:
    def __init__(self):
        self.efficientnet = None
        self.deit = None
        self.models_loaded = False
        self.load_models()
    
    def load_models(self):
        try:
            # Load EfficientNet
            self.efficientnet = get_efficientnet()
            efficientnet_path = os.path.join(CHECKPOINTS_PATH, "efficientnet_best.pth")
            self.efficientnet.load_state_dict(torch.load(efficientnet_path, map_location=device))
            self.efficientnet.eval()
            
            # Load DeiT
            self.deit = get_deit()
            deit_path = os.path.join(CHECKPOINTS_PATH, "deit_best.pth")
            self.deit.load_state_dict(torch.load(deit_path, map_location=device))
            self.deit.eval()
            
            self.models_loaded = True
            print("✅ Models loaded successfully")
        except Exception as e:
            print(f"❌ Error loading models: {str(e)}")
            self.models_loaded = False
    
    @torch.no_grad()
    def predict_efficientnet(self, image):
        if not self.models_loaded:
            return {"error": "Модели не загружены"}
        
        inputs = transform_image(image)
        with autocast():
            outputs = self.efficientnet(inputs)
            probs = torch.nn.functional.softmax(outputs, dim=1)
        
        return self._format_predictions(probs)

    @torch.no_grad()
    def predict_deit(self, image):
        if not self.models_loaded:
            return {"error": "Модели не загружены"}
        
        inputs = transform_image(image)
        with autocast():
            outputs = self.deit(inputs).logits
            probs = torch.nn.functional.softmax(outputs, dim=1)
        
        return self._format_predictions(probs)
    
    @torch.no_grad()
    def predict_ensemble(self, image):
        if not self.models_loaded:
            return {"error": "Модели не загружены"}
        
        inputs = transform_image(image)
        with autocast():
            # Get predictions from both models
            eff_probs = torch.nn.functional.softmax(self.efficientnet(inputs), dim=1)
            deit_probs = torch.nn.functional.softmax(self.deit(inputs).logits, dim=1)
            
            # Ensemble prediction (average probabilities)
            ensemble_probs = (eff_probs + deit_probs) / 2
        
        return self._format_predictions(ensemble_probs)
    
    def _format_predictions(self, probs):
        top5_probs, top5_indices = torch.topk(probs, 5)
        result = {}
        for i in range(5):
            idx = top5_indices[0][i].item()
            label = label_mapping.get(idx, f"Класс {idx}")
            result[label] = float(top5_probs[0][i].item() * 100)
        return result

# Initialize model handler
model_handler = ModelHandler()

# Prediction functions
def predict_efficientnet(image):
    if image is None:
        return "⚠️ Загрузите изображение"
    return model_handler.predict_efficientnet(image)

def predict_deit(image):
    if image is None:
        return "⚠️ Загрузите изображение"
    return model_handler.predict_deit(image)

def predict_ensemble(image):
    if image is None:
        return "⚠️ Загрузите изображение"
    return model_handler.predict_ensemble(image)

# Gradio Interface
def create_individual_tab(model_name, predict_fn):
    with gr.Blocks():
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(label="Загрузите изображение", type="pil")
                predict_btn = gr.Button("Предсказать", variant="primary")
            with gr.Column(scale=1):
                result_output = gr.Label(label="Результаты")
        
        predict_btn.click(
            predict_fn,
            inputs=image_input,
            outputs=result_output
        )
        
        gr.Examples(
            examples=["examples/akiec.jpg", "examples/bcc.jpg", "examples/df.jpg"],
            inputs=image_input,
            label="Примеры из ISIC"
        )

# Create interface
interface = gr.TabbedInterface(
    interface_list=[
        lambda: create_individual_tab("EfficientNet", predict_efficientnet),
        lambda: create_individual_tab("DeiT", predict_deit),
        lambda: create_individual_tab("Ансамблевая модель", predict_ensemble)
    ],
    tab_names=[
        "EfficientNet",
        "DeiT",
        "Ансамблевая модель"
    ],
    title="DermVision Pro",
    description="""
    # Дерматологический классификатор
    Выберите вкладку для использования соответствующей модели:
    - EfficientNet: традиционная CNN модель
    - DeiT: Vision Transformer
    - Ансамблевая модель: комбинация CNN и Vision Transformer
    """,
    theme=gr.themes.Soft(),
    css="""
    .container {max-width: 1200px; margin: auto;}
    .gr-button {font-size: 1.1em; padding: 8px 16px;}
    .gr-textbox {font-size: 1.1em;}
    .gr-column {min-width: 400px;}
    """
)

# Add startup check
def check_models():
    if not model_handler.models_loaded:
        return "⚠️ Предупреждение: Модели не загружены"
    return "✅ Модели готовы к предсказанию"

startup_status = check_models()
print(startup_status)

if __name__ == "__main__":
    print("🚀 Запуск интерфейса...")
    interface.launch()