Update app/fruit_model.py

app/fruit_model.py

@@ -1,6 +1,7 @@
 import os
 import tempfile
 import requests
+import json
 from PIL import Image
 import numpy as np
 import tensorflow as tf
@@ -9,28 +10,33 @@ from tensorflow.keras.models import load_model
 # Optional: Cache download
 os.environ["HF_HOME"] = "/tmp/cache"
 
-# Hugging Face model URL (.keras file)
+# Hugging Face model and labels URLs
 model_url = "https://huggingface.co/jatin1233232/Analiz_Fruit_Model/resolve/main/model.keras"
+labels_url = "https://huggingface.co/jatin1233232/Analiz_Fruit_Model/resolve/main/labels.json"
 
-# Download model to temp file
-def download_model(url):
+# Download file from Hugging Face
+def download_file(url, filename):
     response = requests.get(url)
     if response.status_code == 200:
-        tmp_model_path = os.path.join(tempfile.gettempdir(), "model.keras")
-        with open(tmp_model_path, 'wb') as f:
+        file_path = os.path.join(tempfile.gettempdir(), filename)
+        with open(file_path, 'wb') as f:
             f.write(response.content)
-        return tmp_model_path
+        return file_path
     else:
-        raise Exception("Failed to download model")
+        raise Exception(f"Failed to download: {url}")
 
 # Load the Keras model
-model_path = download_model(model_url)
+model_path = download_file(model_url, "model.keras")
 model = load_model(model_path)
-print("✅ Model loaded successfully.")
+print("✅ Fruit model loaded successfully.")
 
-# Load class labels (optional — if not in model, you need to define them)
-# Example: labels = ['pizza', 'burger', 'idli', ...]
-labels = [f"Class {i}" for i in range(model.output_shape[-1])]  # fallback dummy labels
+# Load labels from labels.json
+labels_path = download_file(labels_url, "labels.json")
+with open(labels_path, 'r') as f:
+    labels_dict = json.load(f)
+
+# Convert to list of labels ordered by index
+labels = [labels_dict[str(i)] for i in range(len(labels_dict))]
 
 # Classification function
 def classify_fruit(image: Image.Image, target_size=(224, 224)):
@@ -45,4 +51,4 @@ def classify_fruit(image: Image.Image, target_size=(224, 224)):
     predicted_label = labels[predicted_index]
     confidence = float(np.max(predictions))
 
-    return predicted_label, confidence
+    return predicted_label, confidence