Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from PIL import Image
 import numpy as np
 from torchvision import models
+from torchvision import transforms
 from transformers import ViTForImageClassification
 from torch import nn
 from torch.cuda.amp import autocast
@@ -24,7 +25,7 @@ label_mapping = {
 }
 
 # Model paths
-CHECKPOINTS_PATH = os.getenv("CHECKPOINTS_PATH", "./")
+CHECKPOINTS_PATH = os.getenv("CHECKPOINTS_PATH", "./checkpoints")
 
 # Model definitions
 def get_efficientnet():
@@ -60,7 +61,7 @@ class ModelHandler:
         self.deit = None
         self.models_loaded = False
         self.load_models()
-    
+
     def load_models(self):
         try:
             # Load EfficientNet
@@ -68,59 +69,56 @@ class ModelHandler:
             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 = {}
@@ -133,7 +131,7 @@ class ModelHandler:
 # Initialize model handler
 model_handler = ModelHandler()
 
-# Prediction functions
+# Prediction wrappers
 def predict_efficientnet(image):
     if image is None:
         return "⚠️ Загрузите изображение"
@@ -149,51 +147,39 @@ def predict_ensemble(image):
         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",
-        "Ансамблевая модель"
-    ]
-)
-
-# Add startup check
-def check_models():
-    if not model_handler.models_loaded:
-        return "⚠️ Предупреждение: Модели не загружены"
-    return "✅ Модели готовы к предсказанию"
-
-startup_status = check_models()
-print(startup_status)
-
+# Create Gradio Blocks with Tabs
+def create_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Диагностика кожных поражений (HAM10K)")
+        status = "✅ Модели готовы к предсказанию" if model_handler.models_loaded else "⚠️ Предупреждение: Модели не загружены"
+        gr.Markdown(f"**Состояние моделей:** {status}")
+
+        with gr.Tabs():
+            with gr.TabItem("EfficientNet"):
+                img = gr.Image(label="Загрузите изображение", type="pil")
+                btn = gr.Button("Предсказать", variant="primary")
+                out = gr.Label(label="Результаты")
+                btn.click(predict_efficientnet, inputs=img, outputs=out)
+                gr.Examples(examples=["examples/akiec.jpg", "examples/bcc.jpg", "examples/df.jpg"], inputs=img)
+
+            with gr.TabItem("DeiT"):
+                img = gr.Image(label="Загрузите изображение", type="pil")
+                btn = gr.Button("Предсказать", variant="primary")
+                out = gr.Label(label="Результаты")
+                btn.click(predict_deit, inputs=img, outputs=out)
+                gr.Examples(examples=["examples/akiec.jpg", "examples/bcc.jpg", "examples/df.jpg"], inputs=img)
+
+            with gr.TabItem("Ансамблевая модель"):
+                img = gr.Image(label="Загрузите изображение", type="pil")
+                btn = gr.Button("Предсказать", variant="primary")
+                out = gr.Label(label="Результаты")
+                btn.click(predict_ensemble, inputs=img, outputs=out)
+                gr.Examples(examples=["examples/akiec.jpg", "examples/bcc.jpg", "examples/df.jpg"], inputs=img)
+
+        return demo
+
+# Launch interface
 if __name__ == "__main__":
+    interface = create_interface()
     print("🚀 Запуск интерфейса...")
-    interface.launch()
+    interface.launch(share=True)