Update app.py

app.py

@@ -3,21 +3,42 @@ import torch
 from PIL import Image
 import gradio as gr
 
-# Загрузка модели
+# Загрузка модели и процессора
 model = AutoModelForImageClassification.from_pretrained("jeemsterri/fish_classification")
 processor = AutoImageProcessor.from_pretrained("jeemsterri/fish_classification")
 
-# Предсказание
+# Функция классификации
 def classify_image(image):
-    inputs = processor(images=image, return_tensors="pt")
-    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()]
-    return {label: float(confidence)}
+    try:
+        # Преобразование изображения и предсказание
+        inputs = processor(images=image, return_tensors="pt")
+        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()]
+        
+        # Формат ответа для API
+        return {
+            "label": label,
+            "confidence": float(confidence),
+            "top_classes": [
+                {"label": model.config.id2label[i], "score": float(probs[0][i])}
+                for i in torch.topk(probs, 3).indices[0]
+            ]
+        }
+    except Exception as e:
+        return {"error": str(e)}
 
-gr.Interface(fn=classify_image, 
-             inputs=gr.Image(type="pil"), 
-             outputs=gr.Label(num_top_classes=3),
-             title="Fish Species Classifier",
-             description="Upload a fish image and get the predicted species.").launch()
+# Gradio Interface
+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."
+)
+
+# Запуск с CORS (если нужно)
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)