Обновление модели и логики обучения

README.md

@@ -9,23 +9,33 @@ app_file: app.py
 pinned: false
 ---
 
-# Классификатор болезней томатов
+# 🍅 Классификатор болезней томатов
 
-Приложение для распознавания болезней томатов с использованием машинного обучения.
+## Описание проекта
+Интерактивное приложение для диагностики болезней томатов с использованием машинного обучения.
 
 ## Возможности
-- Классификация 10 различных состояний листьев томатов
-- Использование SVM-модели с RBF-ядром
-- Высокая точность распознавания
+- Определение болезней по изображению листа томата
+- Поддержка 10 различных классов заболеваний
+- Автоматическое обучение модели при первом запуске
+
+## Технологии
+- Gradio
+- scikit-learn
+- OpenCV
+- PyTorch
 
 ## Классы болезней
-1. Бактериальное пятно
-2. Ранняя пятнистость
-3. Поздняя пятнистость
+1. Бактериальная пятнистость
+2. Ранняя фитофтора
+3. Поздняя фитофтора
 4. Листовая плесень
-5. Септориоз
+5. Пятнистость листьев Септориоз
 6. Паутинный клещ
 7. Целевая пятнистость
-8. Скручивание листьев
+8. Вирус желтой курчавости листьев
 9. Мозаичный вирус
-10. Здоровые листья
+10. Здоровый лист
+
+## Использование
+Загрузите изображение листа томата и получите диагноз с вероятностями.

app.py

@@ -1,29 +1,14 @@
-import os
-import cv2
-import numpy as np
-import torch
-import pkg_resources
-import subprocess
-import sys
-import threading
-import time
-import queue
-
-# Проверка и обновление версии Gradio
-required_gradio = '4.44.1'
-try:
-    current_gradio = pkg_resources.get_distribution('gradio').version
-    if current_gradio != required_gradio:
-        subprocess.check_call([sys.executable, '-m', 'pip', 'install', '--no-cache-dir', f'gradio=={required_gradio}'])
-        print(f'Gradio обновлен до версии {required_gradio}')
-except:
-    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '--no-cache-dir', f'gradio=={required_gradio}'])
-    print(f'Установлен Gradio версии {required_gradio}')
-
 import gradio as gr
+import torch
+import numpy as np
+import cv2
+from sklearn.preprocessing import StandardScaler
+from sklearn.svm import SVC
+import logging
+from train import download_and_prepare_dataset, load_images_and_labels, train_and_evaluate_model
 
 # Список классов болезней
-TOMATO_CLASSES = [
+DISEASE_CLASSES = [
     'Tomato___Bacterial_spot', 
     'Tomato___Early_blight', 
     'Tomato___Late_blight', 
@@ -36,28 +21,31 @@ TOMATO_CLASSES = [
     'Tomato___healthy'
 ]
 
-# Состояние обучения
-MODEL_TRAINING_STATUS = queue.Queue()
-MODEL_TRAINING_STATUS.put("Модель не обучена")
-
 def load_model():
     """Загрузка обученной модели"""
     try:
         model_path = 'tomato_disease_classifier.pth'
         
-        if not os.path.exists(model_path):
-            print(f"Модель не найдена по пути: {os.path.abspath(model_path)}")
-            return None, None
+        # Если модель не существует, обучаем
+        if not torch.os.path.exists(model_path):
+            download_and_prepare_dataset()
+            X, y = load_images_and_labels()
+            train_and_evaluate_model(X, y)
         
+        # Загрузка модели
         model_data = torch.load(model_path)
-        if 'classifier' not in model_data or 'scaler' not in model_data:
-            print("Некорректная структура модели: отсутствуют необходимые компоненты")
-            return None, None
         
-        print("Модель успешно загружена")
-        return model_data['classifier'], model_data['scaler']
+        # Создание scaler
+        scaler = StandardScaler()
+        scaler.mean_ = model_data['mean']
+        scaler.scale_ = model_data['std']
+        
+        classifier = model_data['classifier']
+        
+        return scaler, classifier
+    
     except Exception as e:
-        print(f"Ошибка загрузки модели: {str(e)}")
+        print(f"Ошибка загрузки модели: {e}")
         return None, None
 
 def preprocess_image(image):
@@ -74,12 +62,12 @@ def preprocess_image(image):
 def predict_disease(image):
     """Предсказание болезни томата"""
     if image is None:
-        return "Пожалуйста, загрузите изображение листа томата"
+        return "Пожалуйста, загрузите изображение"
     
     # Загрузка модели
-    classifier, scaler = load_model()
-    if classifier is None or scaler is None:
-        return "Модель не найдена. Пожалуйста, убедитесь, что файл 'tomato_disease_classifier.pth' находится в корневой папке проекта."
+    scaler, classifier = load_model()
+    if scaler is None or classifier is None:
+        return "Ошибка загрузки модели. Возможно, нужно сначала обучить модель."
     
     # Предобработка изображения
     processed_image = preprocess_image(image)
@@ -96,61 +84,20 @@ def predict_disease(image):
     # Формирование результата
     result = f"Обнаружено: {prediction[0]}\n\n"
     result += "Вероятности:\n"
-    for disease, prob in zip(TOMATO_CLASSES, probabilities):
+    for disease, prob in zip(DISEASE_CLASSES, probabilities):
         result += f"{disease}: {prob*100:.2f}%\n"
     
     return result
 
 # Создание Gradio интерфейса
-demo = gr.Blocks(title="Диагностика болезней томатов")
+iface = gr.Interface(
+    fn=predict_disease,
+    inputs=gr.Image(type="numpy", label="Загрузите изображение листа томата"),
+    outputs=gr.Textbox(label="Результат диагностики"),
+    title="🍅 Диагностика болезней томатов",
+    description="Загрузите изображение листа томата для определения заболевания"
+)
 
-with demo:
-    gr.Markdown("""
-    # Диагностика болезней томатов
-    За��рузите изображение листа томата для определения заболевания
-    """)
-    
-    status_text = gr.Textbox(label="Статус модели", interactive=False)
-    
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="numpy", label="Загрузите изображение листа томата")
-            predict_btn = gr.Button("Определить болезнь")
-        
-        with gr.Column():
-            output_text = gr.Textbox(label="Результат диагностики")
-    
-    predict_btn.click(
-        fn=predict_disease,
-        inputs=input_image,
-        outputs=output_text
-    )
-    
-    # Периодическое обновление статуса
-    demo.load(fn=lambda: MODEL_TRAINING_STATUS.queue[0] if not MODEL_TRAINING_STATUS.empty() else "Модель не обучена", inputs=None, outputs=status_text, every=5)
-
-def train_model_in_background():
-    """Запуск обучения модели в фоновом режиме"""
-    try:
-        MODEL_TRAINING_STATUS.queue.clear()
-        MODEL_TRAINING_STATUS.put("Начало обучения модели...")
-        print("Начинаем обучение модели...")
-        from train import main as train_main
-        train_main()
-        MODEL_TRAINING_STATUS.queue.clear()
-        MODEL_TRAINING_STATUS.put("Модель успешно обучена!")
-        print("Модель успешно обучена!")
-    except Exception as e:
-        MODEL_TRAINING_STATUS.queue.clear()
-        MODEL_TRAINING_STATUS.put(f"Ошибка при обучении: {str(e)}")
-        print(f"Ошибка при обучении модели: {str(e)}")
-
-# Запуск интерфейса
+# Запуск приложения
 if __name__ == "__main__":
-    # Проверка наличия модели
-    if not os.path.exists('tomato_disease_classifier.pth'):
-        print("Модель не найдена, запускаем обучение в фоновом режиме")
-        training_thread = threading.Thread(target=train_model_in_background)
-        training_thread.start()
-    
-    demo.launch()
+    iface.launch(server_name="0.0.0.0", server_port=7860)

hf_app.py

@@ -7,6 +7,8 @@ from fastapi import FastAPI, File, UploadFile
 from sklearn.preprocessing import StandardScaler
 from sklearn.svm import SVC
 import logging
+from train import download_and_prepare_dataset, load_images_and_labels, train_and_evaluate_model, cleanup
+import threading
 
 # Настройка логгера
 logger = logging.getLogger(__name__)
@@ -51,7 +53,8 @@ def load_model():
         model_path = next((path for path in model_paths if os.path.exists(path)), None)
         
         if model_path is None:
-            logger.error("Модель не найдена ни в одном из путей")
+            logger.info("Модель не найдена, запускаем обучение в фоновом режиме")
+            threading.Thread(target=train_model).start()
             return None, None
         
         logger.info(f"Загрузка модели из: {model_path}")
@@ -122,5 +125,17 @@ def read_root():
     return {"status": "Tomato Disease Classifier is running"}
 
 # Запуск Gradio
+def train_model():
+    try:
+        logger.info("Начинаем обучение модели...")
+        download_and_prepare_dataset()
+        X, y = load_images_and_labels()
+        train_and_evaluate_model(X, y)
+        logger.info("Модель успешно обучена!")
+    except Exception as e:
+        logger.error(f"Ошибка при обучении модели: {e}")
+    finally:
+        cleanup()
+
 if __name__ == "__main__":
     iface.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -1,18 +1,18 @@
 # Gradio и его зависимости
 gradio==4.44.1
-fastapi==0.104.1
-uvicorn[standard]==0.27.0
+fastapi
+uvicorn[standard]
 
 # Основные библиотеки для ML
-torch==2.1.2
-numpy==1.24.3
-opencv-python-headless==4.8.1.78
-scikit-learn==1.3.2
-matplotlib==3.7.2
-seaborn==0.12.2
+torch
+numpy
+opencv-python-headless
+scikit-learn
+matplotlib
+seaborn
 
 # Вспомогательные библиотеки
-python-multipart==0.0.9
+python-multipart
 requests==2.31.0
 scipy==1.12.0
 tqdm==4.66.2

train.py

@@ -152,7 +152,8 @@ def train_and_evaluate_model(X, y):
     model_save_path = os.path.join(os.getcwd(), 'tomato_disease_classifier.pth')
     torch.save({
         'classifier': classifier,
-        'scaler': scaler,
+        'mean': scaler.mean_,
+        'std': scaler.scale_,
         'classes': TOMATO_CLASSES
     }, model_save_path)