Update app.py

app.py

@@ -1,100 +1,73 @@
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-import torch
-from PIL import Image
 import gradio as gr
+import requests
+from PIL import Image
+import numpy as np
 import logging
-from functools import lru_cache
-import os
 
-# Настройка логирования
+# Configure logging
 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}")
+# Hugging Face API settings
+HF_API_URL = "https://api-inference.huggingface.co/models/jeemsterri/fish_classification"
+HF_API_KEY = "your_huggingface_api_key"  # Replace with your key
 
-# Кэширование загрузки модели
-@lru_cache(maxsize=1)
-def load_model():
-    logger.info("Loading model and processor...")
+def classify_fish(image: Image.Image) -> dict:
+    """
+    Classify a fish image using Hugging Face API or fallback to MobileNet.
+    Args:
+        image: PIL Image object.
+    Returns:
+        Dict with predictions or error message.
+    """
     try:
-        model = AutoModelForImageClassification.from_pretrained(
-            "jeemsterri/fish_classification"
-        ).to(device)
-        processor = AutoImageProcessor.from_pretrained("jeemsterri/fish_classification", use_fast=True)  # Используем быстрый процессор
-        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
+        # Convert image to bytes for API
+        img_bytes = image.tobytes()
 
-def classify_image(image):
-    try:
-        # Конвертация numpy array в PIL Image
-        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)
+        # Try Hugging Face API first
+        headers = {"Authorization": f"Bearer {HF_API_KEY}"}
+        response = requests.post(HF_API_URL, headers=headers, data=img_bytes)
         
-        # Обработка результатов
-        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-        confidence, predicted_class = torch.max(probs, dim=1)
-        top_classes = torch.topk(probs, 3)
+        if response.status_code == 200:
+            predictions = response.json()
+            logger.info(f"API response: {predictions}")
+            return {"source": "Hugging Face", "predictions": predictions}
         
-        # Формирование результата
-        result = {
-            "label": model.config.id2label[predicted_class.item()],
-            "confidence": float(confidence),
-            "top_classes": [
-                {
-                    "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
+        # Fallback to MobileNet if API fails
+        logger.warning(f"API failed (status {response.status_code}), using fallback...")
+        import tensorflow as tf
+        import tensorflow_hub as hub
+
+        # Load MobileNet
+        model = tf.keras.Sequential([
+            hub.KerasLayer("https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4")
+        ])
+        image = image.resize((224, 224))  # MobileNet expects 224x224
+        image_array = np.array(image) / 255.0
+        image_array = np.expand_dims(image_array, axis=0)
+
+        predictions = model.predict(image_array)
+        top_prediction = tf.keras.applications.mobilenet_v2.decode_predictions(predictions, top=1)[0][0]
         
-    except Exception as e:
-        error_msg = f"Classification error: {str(e)}"
-        logger.error(error_msg)
         return {
-            "error": error_msg,
-            "available_labels": list(model.config.id2label.values())[:10] + ["..."]
+            "source": "MobileNet (Fallback)",
+            "predictions": [{"label": top_prediction[1], "score": float(top_prediction[2])}]
         }
 
-# Создание интерфейса Gradio
-iface = gr.Interface(
-    fn=classify_image,
+    except Exception as e:
+        logger.error(f"Classification error: {str(e)}")
+        return {"error": str(e)}
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=classify_fish,
     inputs=gr.Image(type="pil", label="Upload Fish Image"),
-    outputs=gr.JSON(label="Classification Results"),
-    title="🐟 Fish Species Classifier",
-    description="Upload an image of a fish to identify its species",
-    examples=None,  # Убраны примеры, так как файлы отсутствуют
-    flagging_mode="never",  # Заменяет устаревший allow_flagging
-    cache_examples=False  # Отключаем кэширование примеров
+    outputs=gr.JSON(label="Prediction Results"),
+    title="🐟 Fish Classifier",
+    description="Upload an image of a fish to see the predicted class probabilities.",
+    examples=["salmon.jpg", "tuna.jpg"],  # Add example images
+    theme="soft"
 )
 
-# Запуск приложения
 if __name__ == "__main__":
-    iface.launch(
-        server_name="0.0.0.0",
-        server_port=7860
-    )
+    interface.launch(server_name="0.0.0.0", server_port=7860)