Update app.py

app.py

@@ -2,43 +2,109 @@ from transformers import AutoImageProcessor, AutoModelForImageClassification
 import torch
 from PIL import Image
 import gradio as gr
+import logging
+from functools import lru_cache
 
-# Загрузка модели и процессора
-model = AutoModelForImageClassification.from_pretrained("jeemsterri/fish_classification")
-processor = AutoImageProcessor.from_pretrained("jeemsterri/fish_classification")
+# Настройка логирования
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Проверка доступности GPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger.info(f"Using device: {device}")
+
+# Кэширование загрузки модели для ускорения последующих запросов
+@lru_cache(maxsize=1)
+def load_model():
+    logger.info("Loading model and processor...")
+    try:
+        model = AutoModelForImageClassification.from_pretrained(
+            "jeemsterri/fish_classification"
+        ).to(device)
+        processor = AutoImageProcessor.from_pretrained("jeemsterri/fish_classification")
+        logger.info("Model loaded successfully")
+        return model, processor
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        raise
+
+# Загрузка модели при старте
+try:
+    model, processor = load_model()
+except Exception as e:
+    logger.error(f"Failed to load model: {str(e)}")
+    raise
 
-# Функция классификации
 def classify_image(image):
     try:
-        # Преобразование изображения и предсказание
-        inputs = processor(images=image, return_tensors="pt")
+        # Проверка входного изображения
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+            
+        logger.info("Processing image...")
+        
+        # Преобразование изображения
+        inputs = processor(images=image, return_tensors="pt").to(device)
+        
+        # Предсказание
         with torch.no_grad():
             outputs = model(**inputs)
+        
+        # Обработка результатов
         probs = torch.nn.functional.softmax(outputs.logits, dim=1)
         confidence, predicted_class = torch.max(probs, dim=1)
-        label = model.config.id2label[predicted_class.item()]
+        top_classes = torch.topk(probs, 3)
         
-        # Формат ответа для API
-        return {
-            "label": label,
+        # Формирование результата
+        result = {
+            "label": model.config.id2label[predicted_class.item()],
             "confidence": float(confidence),
             "top_classes": [
-                {"label": model.config.id2label[i], "score": float(probs[0][i])}
-                for i in torch.topk(probs, 3).indices[0]
+                {
+                    "label": model.config.id2label[i.item()],
+                    "score": float(probs[0][i])
+                }
+                for i in top_classes.indices[0]
             ]
         }
+        
+        logger.info(f"Prediction result: {result}")
+        return result
+        
     except Exception as e:
-        return {"error": str(e)}
+        error_msg = f"Classification error: {str(e)}"
+        logger.error(error_msg)
+        return {
+            "error": error_msg,
+            "available_labels": list(model.config.id2label.values())[:10] + ["..."]
+        }
 
-# Gradio Interface
+# Создание интерфейса Gradio с улучшенным UI
 iface = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.JSON(),
-    title="Fish Species Classifier",
-    description="Upload a fish image to classify its species."
+    inputs=gr.Image(
+        type="pil",
+        label="Upload Fish Image",
+        sources=["upload", "webcam", "clipboard"]
+    ),
+    outputs=gr.JSON(
+        label="Classification Results"
+    ),
+    title="🐟 Fish Species Classifier",
+    description="Upload an image of a fish to identify its species using AI",
+    examples=[
+        ["salmon.jpg"],
+        ["trout.jpg"]
+    ],
+    allow_flagging="never",
+    theme=gr.themes.Soft()
 )
 
-# Запуск с CORS (если нужно)
+# Конфигурация запуска
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        enable_queue=True,
+        share=False
+    )