Update app.py

app.py

@@ -5,132 +5,109 @@ from PIL import Image
 import os
 from typing import Tuple, List, Dict
 
-# Устройство
+# Единоразовое определение устройства
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Загрузка модели
 def load_model() -> Tuple[torch.nn.Module, List[str]]:
+    """
+    Загружает модель и список меток классов.
+    Returns:
+        model: Загруженная модель PyTorch.
+        labels: Список меток классов.
+    """
     model_path = "skinconvnext_scripted.pt"
     labels_path = "labels.txt"
+    
     if not os.path.exists(model_path):
         raise FileNotFoundError(f"Model not found: {model_path}")
     if not os.path.exists(labels_path):
-        raise FileNotFoundError("labels.txt not found.")
+        raise FileNotFoundError("File labels.txt not found.")
+    
     model = torch.jit.load(model_path, map_location=device)
-    model.eval().to(device)
+    model.eval()
+    model.to(device)  # Перемещаем модель на устройство сразу
+    
     with open(labels_path, "r") as f:
         labels = f.read().splitlines()
+    
     return model, labels
 
 model, labels = load_model()
 
-# Преобразования
+# Определение преобразований изображения (создаётся один раз)
 preprocess = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
 ])
 
-# Функция предсказаний
 def predict(image: Image.Image) -> Dict[str, float]:
+    """
+    Выполняет предсказание для переданного изображения.
+    Args:
+        image (PIL.Image): Входное изображение.
+    
+    Returns:
+        Dict[str, float]: Словарь, где ключи – имена классов, а значения – вероятности.
+    """
     try:
-        img = image.convert("RGB")
-        tensor = preprocess(img).unsqueeze(0).to(device)
+        # Приводим изображение к RGB и преобразуем
+        image = image.convert("RGB")
+        image_tensor = preprocess(image).unsqueeze(0).to(device)
+        
         with torch.no_grad():
-            out = model(tensor)
-            probs = torch.nn.functional.softmax(out[0], dim=0)
-        preds = {lbl: float(p) for lbl,p in zip(labels, probs)}
-        return dict(sorted(preds.items(), key=lambda x: x[1], reverse=True))
+            output = model(image_tensor)
+            # Предполагаем, что output имеет размерность [1, N]
+            scores = torch.nn.functional.softmax(output[0], dim=0)
+        
+        # Формируем словарь с предсказаниями
+        predictions = {label: float(score) for label, score in zip(labels, scores)}
+        sorted_predictions = dict(sorted(predictions.items(), key=lambda item: item[1], reverse=True))
+        return sorted_predictions
     except Exception as e:
         return {"error": str(e)}
 
-# Примеры
-examples = [["example1.jpg"], ["example2.jpg"]]
+# Описание интерфейса Gradio
+title = "🔥 Skin-AI"
+description = (
+    "🔬 **Skin-AI — AI-Powered Skin Disease Classification**\n\n"
+    "Проект использует глубокую модель для классификации заболеваний кожи по изображению.\n\n"
+    "### 🚀 Как использовать:\n\n"
+    "1️⃣ Загрузите или сфотографируйте поражённый участок кожи.\n\n"
+    "2️⃣ Нажмите кнопку 'Submit'.\n\n"
+    "3️⃣ Приложение покажет вероятности для возможных заболеваний.\n\n"
+    "⚠️ **Внимание!** Результаты предоставлены в ознакомительных целях и не являются медицинской диагностикой.\n\n"
+    "### 🛠 Используемые технологии:\n"
+    "- PyTorch\n"
+    "- Gradio\n"
+    "- Hugging Face Spaces\n\n"
+    "🔗 Исходный код: [Hugging Face](https://huggingface.co/Eraly-ml/Skin-AI)"
+)
 
-# Переводы
-translations = {
-    "Русский": {
-        "title": "# 🔥 Skin-AI",
-        "description": (
-            "🔬 **Skin-AI — AI для классификации кожных заболеваний**\n\n"
-            "1️⃣ Загрузите или выберите пример ниже.\n"
-            "2️⃣ Нажмите кнопку «Предсказать».\n"
-            "⚠️ **Не является медицинской диагностикой.**"
-        ),
-        "upload_label": "Загрузить изображение",
-        "submit": "Предсказать",
-        "result": "Результат"
-    },
-    "Қазақша": {
-        "title": "# 🔥 Skin-AI",
-        "description": (
-            "🔬 **Skin-AI — Терi ауруларын анықтайтын ИИ**\n\n"
-            "1️⃣ Суретті жүктеңіз не төмендегі мысалдарды таңдаңыз.\n"
-            "2️⃣ «Болжам жасау» батырмасын басыңыз.\n"
-            "⚠️ **Бұл медициналық диагноз емес.**"
-        ),
-        "upload_label": "Суретті жүктеу",
-        "submit": "Болжам жасау",
-        "result": "Нәтиже"
-    },
-    "English": {
-        "title": "# 🔥 Skin-AI",
-        "description": (
-            "🔬 **Skin-AI — AI-Powered Skin Disease Classification**\n\n"
-            "1️⃣ Upload an image or choose an example below.\n"
-            "2️⃣ Click the “Submit” button.\n"
-            "⚠️ **This is not a medical diagnosis.**"
-        ),
-        "upload_label": "Upload Image",
-        "submit": "Submit",
-        "result": "Prediction"
-    }
-}
+# Примеры изображений
+examples = [
+    ["example1.jpg"],
+    ["example2.jpg"]
+]
 
-# Функция переключения языка
-def change_language(lang):
-    tr = translations[lang]
-    return (
-        gr.update(value=tr["title"]),
-        gr.update(value=tr["description"]),
-        gr.update(label=tr["upload_label"]),
-        gr.update(value=tr["submit"]),
-        gr.update(label=tr["result"])
-    )
+def update_submit_state(image):
+    return gr.update(interactive=image is not None)
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # Селектор языка
-    lang = gr.Dropdown(
-        choices=["Русский","Қазақша","English"],
-        value="Русский",
-        label="🌐 Язык / Language / Тiлдер"
-    )
-
-    # Заголовок и описание
-    title_md = gr.Markdown(translations["Русский"]["title"])
-    desc_md  = gr.Markdown(translations["Русский"]["description"])
-
-    # Основные компоненты
-    image_input   = gr.Image(type="pil", label=translations["Русский"]["upload_label"])
-    submit_button = gr.Button(translations["Русский"]["submit"])
-    output_label  = gr.Label(num_top_classes=3, label=translations["Русский"]["result"])
-
-    # Примеры
-    gr.Examples(
-        examples=examples,
-        inputs=image_input,
-        outputs=output_label,
-        fn=predict,
-        cache_examples=True
-    )
-
-    # Связи
-    lang.change(
-        fn=change_language,
-        inputs=lang,
-        outputs=[title_md, desc_md, image_input, submit_button, output_label]
-    )
+with gr.Blocks(theme=gr.themes.Soft()) as interface:
+    gr.Markdown(title)
+    gr.Markdown(description)
+    
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Image")
+    
+    submit_button = gr.Button("Submit", interactive=False)
+    output_label = gr.Label(num_top_classes=3, label="Prediction")
+    
+    image_input.change(fn=update_submit_state, inputs=image_input, outputs=submit_button)
     submit_button.click(fn=predict, inputs=image_input, outputs=output_label)
-
+    
+    gr.Examples(examples, inputs=image_input)
+    
 if __name__ == "__main__":
-    demo.launch()
+    interface.launch()