Update app.py

app.py

@@ -1,11 +1,36 @@
 import gradio as gr
 import requests
-from transformers import pipeline
 from PIL import Image
+import torch
+from torchvision import models
+from huggingface_hub import hf_hub_download
+from torchvision import transforms
 
-# Menggunakan pipeline dari Hugging Face
-model_name = "ahmadalfian/fruits_vegetables_classifier"  # Model klasifikasi
-classifier = pipeline("image-classification", model=model_name)
+# Mengunduh dan mempersiapkan model
+model_path = hf_hub_download(repo_id="ahmadalfian/fruits_vegetables_classifier", filename="resnet50_finetuned.pth")
+model = models.resnet50(pretrained=False)
+num_classes = 36
+model.fc = torch.nn.Linear(in_features=2048, out_features=num_classes)
+model.load_state_dict(torch.load(model_path))
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
+
+# Transformasi untuk pra-pemrosesan gambar
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Fungsi untuk memprediksi kelas
+def predict(image):
+    image = image.convert("RGB")
+    image_tensor = preprocess(image).unsqueeze(0).to(device)  # Tambahkan dimensi batch dan pindahkan ke device
+    with torch.no_grad():
+        outputs = model(image_tensor)
+    predictions = outputs.argmax(dim=1)
+    return predictions.item()
 
 # Fungsi untuk mengambil informasi nutrisi
 def get_nutritional_info(food):
@@ -13,8 +38,8 @@ def get_nutritional_info(food):
     url = "https://api.nal.usda.gov/fdc/v1/foods/search"
     
     params = {
-        "query": food,  # Nama makanan yang diprediksi
-        "pageSize": 5,  # Ambil lima hasil
+        "query": food,
+        "pageSize": 5,
         "api_key": api_key
     }
     
@@ -22,7 +47,6 @@ def get_nutritional_info(food):
     data = response.json()
 
     if "foods" in data and len(data["foods"]) > 0:
-        # Inisialisasi total untuk nutrisi yang diinginkan
         nutrients_totals = {
             "Energy": 0,
             "Carbohydrate, by difference": 0,
@@ -31,30 +55,34 @@ def get_nutritional_info(food):
         }
         item_count = len(data["foods"])
 
-        # Iterasi melalui setiap makanan
         for food in data["foods"]:
             for nutrient in food['foodNutrients']:
                 nutrient_name = nutrient['nutrientName']
                 nutrient_value = nutrient['value']
 
-                # Cek apakah nutrisi termasuk yang diinginkan
                 if nutrient_name in nutrients_totals:
                     nutrients_totals[nutrient_name] += nutrient_value
 
-        # Menghitung rata-rata nilai nutrisi
         average_nutrients = {name: total / item_count for name, total in nutrients_totals.items()}
 
         return average_nutrients
     else:
         return None
 
-# Fungsi utama untuk klasifikasi dan informasi nutrisi
+# Fungsi utama Gradio
 def classify_and_get_nutrition(image):
-    # Prediksi menggunakan pipeline
-    predictions = classifier(image)
-    predicted_label = predictions[0]['label'].lower()  # Ambil label prediksi
+    predicted_class_idx = predict(image)
+    class_labels = [
+        'apple', 'banana', 'beetroot', 'bell pepper', 'cabbage', 'capsicum',
+        'carrot', 'cauliflower', 'chilli pepper', 'corn', 'cucumber', 
+        'eggplant', 'garlic', 'ginger', 'grapes', 'jalepeno', 'kiwi', 
+        'lemon', 'lettuce', 'mango', 'onion', 'orange', 'paprika', 
+        'pear', 'peas', 'pineapple', 'pomegranate', 'potato', 'raddish', 
+        'soy beans', 'spinach', 'sweetcorn', 'sweetpotato', 'tomato', 
+        'turnip', 'watermelon'
+    ]
+    predicted_label = class_labels[predicted_class_idx]
     
-    # Ambil informasi nutrisi
     nutrisi = get_nutritional_info(predicted_label)
     
     if nutrisi: