Update app.py

app.py

@@ -2,8 +2,9 @@ import gradio as gr
 import torch
 from PIL import Image
 import numpy as np
-from torchvision import models
-from torchvision import transforms
+from torchvision import models, transforms
+from torchvision.datasets import ImageFolder
+from torch.utils.data import DataLoader
 from transformers import ViTForImageClassification
 from torch import nn
 from torch.cuda.amp import autocast
@@ -25,35 +26,39 @@ label_mapping = {
     6: "Сосудистые поражения"
 }
 
-# Model paths
-CHECKPOINTS_PATH = os.getenv("CHECKPOINTS_PATH", "./")
+# Paths and hyperparams
+CHECKPOINTS_PATH = os.getenv("CHECKPOINTS_PATH", "./checkpoints")
+SUBMISSIONS_PATH = os.getenv("SUBMISSIONS_PATH", "./submissions")
+FT_BATCH = 32
+FT_EPOCHS = 1  # adjust as needed
+LR = 1e-4
+
+os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
+os.makedirs(SUBMISSIONS_PATH, exist_ok=True)
 
 # Model definitions
 def get_efficientnet():
     model = models.efficientnet_v2_s(weights="IMAGENET1K_V1")
-    model.classifier[1] = nn.Linear(1280, 7)
+    model.classifier[1] = nn.Linear(1280, len(label_mapping))
     return model.to(device)
 
 def get_deit():
     model = ViTForImageClassification.from_pretrained(
         'facebook/deit-base-patch16-224',
-        num_labels=7,
+        num_labels=len(label_mapping),
         ignore_mismatched_sizes=True
     )
     return model.to(device)
 
 # Transforms
+train_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])
+])
+
 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)
+    return train_transform(image).unsqueeze(0).to(device)
 
 # Model Handler
 class ModelHandler:
@@ -65,13 +70,11 @@ class ModelHandler:
 
     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))
+            eff_path = os.path.join(CHECKPOINTS_PATH, "efficientnet_best.pth")
+            self.efficientnet.load_state_dict(torch.load(eff_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))
@@ -80,79 +83,111 @@ class ModelHandler:
             self.models_loaded = True
             print("✅ Models loaded successfully")
         except Exception as e:
-            print(f"❌ Error loading models: {str(e)}")
+            print(f"❌ Error loading models: {e}")
             self.models_loaded = False
 
     @torch.no_grad()
-    def predict_efficientnet(self, image):
-        if not self.models_loaded:
-            return {"error": "Модели не загружены"}
-
-        inputs = transform_image(image)
-        # Handle autocast based on device
-        ctx = autocast() if device.type == 'cuda' else nullcontext()
-        with ctx:
-            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):
+    def predict(self, image, use='efficientnet'):
         if not self.models_loaded:
             return {"error": "Модели не загружены"}
-
         inputs = transform_image(image)
         ctx = autocast() if device.type == 'cuda' else nullcontext()
         with ctx:
-            outputs = self.deit(pixel_values=inputs).logits  # Corrected parameter
-            probs = torch.nn.functional.softmax(outputs, dim=1)
-
+            if use == 'efficientnet':
+                logits = self.efficientnet(inputs)
+            elif use == 'deit':
+                logits = self.deit(pixel_values=inputs).logits
+            else:
+                logits = (self.efficientnet(inputs) + self.deit(pixel_values=inputs).logits) / 2
+            probs = torch.nn.functional.softmax(logits, dim=1)
         return self._format_predictions(probs)
 
-    @torch.no_grad()
-    def predict_ensemble(self, image):
-        if not self.models_loaded:
-            return {"error": "Модели не загружены"}
+    def _format_predictions(self, probs):
+        top5_probs, top5_inds = torch.topk(probs, 5)
+        return {label_mapping[i.item()]: float(top5_probs[0][k].item())
+                for k, i in enumerate(top5_inds[0])}
 
-        inputs = transform_image(image)
-        ctx = autocast() if device.type == 'cuda' else nullcontext()
-        with ctx:
-            eff_probs = torch.nn.functional.softmax(self.efficientnet(inputs), dim=1)
-            deit_probs = torch.nn.functional.softmax(self.deit(pixel_values=inputs).logits, dim=1)
-            ensemble_probs = (eff_probs + deit_probs) / 2
-
-        return self._format_predictions(ensemble_probs)
-
-    def _format_predictions(self, probs):  # Corrected indentation
-        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())
-        return result
-
-# Initialize model handler
+# Initialize handler
 model_handler = ModelHandler()
 
-# Prediction wrappers
 def predict_efficientnet(image):
-    if image is None:
-        return "⚠️ Загрузите изображение"
-    return model_handler.predict_efficientnet(image)
+    return "⚠️ Загрузите изображение" if image is None else model_handler.predict(image, 'efficientnet')
 
 def predict_deit(image):
-    if image is None:
-        return "⚠️ Загрузите изображение"
-    return model_handler.predict_deit(image)
+    return "⚠️ Загрузите изображение" if image is None else model_handler.predict(image, 'deit')
 
 def predict_ensemble(image):
-    if image is None:
-        return "⚠️ Загрузите изображение"
-    return model_handler.predict_ensemble(image)
-
-# Create Gradio Blocks with Tabs
+    return "⚠️ Загрузите изображение" if image is None else model_handler.predict(image, 'ensemble')
+
+# Finetuning logic
+
+def finetune_models():
+    # Prepare dataset
+    dataset = ImageFolder(SUBMISSIONS_PATH, transform=train_transform)
+    loader = DataLoader(dataset, batch_size=8, shuffle=True)
+
+    # Finetune EfficientNet
+    eff = get_efficientnet()
+    eff.load_state_dict(torch.load(os.path.join(CHECKPOINTS_PATH, "efficientnet_best.pth"), map_location=device))
+    eff.train()
+    optimizer = torch.optim.Adam(eff.parameters(), lr=LR)
+    criterion = nn.CrossEntropyLoss()
+    for epoch in range(FT_EPOCHS):
+        for imgs, lbls in loader:
+            imgs, lbls = imgs.to(device), lbls.to(device)
+            optimizer.zero_grad()
+            outputs = eff(imgs)
+            loss = criterion(outputs, lbls)
+            loss.backward()
+            optimizer.step()
+    torch.save(eff.state_dict(), os.path.join(CHECKPOINTS_PATH, "efficientnet_best.pth"))
+
+    # Finetune DeiT
+    dt = get_deit()
+    dt.load_state_dict(torch.load(os.path.join(CHECKPOINTS_PATH, "deit_best.pth"), map_location=device))
+    dt.train()
+    optimizer = torch.optim.Adam(dt.parameters(), lr=LR)
+    for epoch in range(FT_EPOCHS):
+        for imgs, lbls in loader:
+            imgs, lbls = imgs.to(device), lbls.to(device)
+            optimizer.zero_grad()
+            outputs = dt(pixel_values=imgs).logits
+            loss = criterion(outputs, lbls)
+            loss.backward()
+            optimizer.step()
+    torch.save(dt.state_dict(), os.path.join(CHECKPOINTS_PATH, "deit_best.pth"))
+
+    # Reload into handler
+    model_handler.load_models()
+    print("🔄 Models fine-tuned and reloaded")
+
+
+def handle_submission(image, label):
+    if image is None or label is None:
+        return "⚠️ Загрузите изображение и выберите метку"
+    # Save image under label folder
+    lbl_dir = os.path.join(SUBMISSIONS_PATH, str(label))
+    os.makedirs(lbl_dir, exist_ok=True)
+    idx = len([f for f in os.listdir(lbl_dir) if f.endswith(('.png','.jpg'))]) + 1
+    path = os.path.join(lbl_dir, f"{label}_{idx}.png")
+    image.save(path)
+
+    # Count total submissions
+    total = sum(len(files) for _, _, files in os.walk(SUBMISSIONS_PATH))
+    rem = FT_BATCH - (total % FT_BATCH)
+    if rem == FT_BATCH:
+        rem = 0  # just reached batch multiple
+    # Trigger finetune if batch complete
+    if total % FT_BATCH == 0:
+        finetune_models()
+        # Clear submissions
+        for root, _, files in os.walk(SUBMISSIONS_PATH):
+            for f in files:
+                os.remove(os.path.join(root, f))
+
+    return f"Осталось {rem} изображений до следующей тонкой настройки"
+
+# Create Gradio interface
 def create_interface():
     with gr.Blocks() as demo:
         gr.Markdown("# Диагностика кожных поражений (HAM10K)")
@@ -161,30 +196,27 @@ def create_interface():
 
         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)
+                img, out = gr.Image(type="pil", label="Загрузите изображение"), gr.Label(label="Результаты")
+                gr.Button("Предсказать").click(predict_efficientnet, inputs=img, outputs=out)
 
             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)
+                img, out = gr.Image(type="pil", label="Загрузите изображение"), gr.Label(label="Результаты")
+                gr.Button("Предсказать").click(predict_deit, inputs=img, outputs=out)
 
             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)
+                img, out = gr.Image(type="pil", label="Загрузите изображение"), gr.Label(label="Результаты")
+                gr.Button("Предсказать").click(predict_ensemble, inputs=img, outputs=out)
+
+            with gr.TabItem("Submit for Finetuning"):
+                sub_img = gr.Image(type="pil", label="Изображение для тонкой настройки")
+                sub_lbl = gr.Dropdown(choices=list(label_mapping.values()), label="Выберите метку")
+                sub_btn = gr.Button("Отправить")
+                sub_out = gr.Textbox(label="Статус")
+                sub_btn.click(handle_submission, inputs=[sub_img, sub_lbl], outputs=sub_out)
 
         return demo
 
-# Launch interface
 if __name__ == "__main__":
     interface = create_interface()
     print("🚀 Запуск интерфейса...")
-    interface.launch(server_port=7860)  # Explicitly set port if needed
+    interface.launch(server_port=7860)