Update app.py

app.py

@@ -1,34 +1,34 @@
 import gradio as gr
 import requests
-from PIL import Image
-import torch
-from torchvision import models
 from huggingface_hub import hf_hub_download
-from torchvision import transforms
+import torch
+from PIL import Image
 
-# Mengunduh dan mempersiapkan model
+# Mengunduh model dari Hugging Face
 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]),
-])
+# Menginisialisasi model
+class FruitVegetableClassifier(torch.nn.Module):
+    def __init__(self):
+        super(FruitVegetableClassifier, self).__init__()
+        self.model = torch.load(model_path)
+        self.model.eval()
+
+    def forward(self, x):
+        return self.model(x)
+
+model = FruitVegetableClassifier()
 
 # Fungsi untuk memprediksi kelas
 def predict(image):
+    # Proses gambar di sini (ubah ukuran, normalisasi, dll.)
     image = image.convert("RGB")
-    image_tensor = preprocess(image).unsqueeze(0).to(device)  # Tambahkan dimensi batch dan pindahkan ke device
+    image = image.resize((224, 224))  # Resize ke ukuran yang diharapkan model
+    image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float().unsqueeze(0)  # Ubah menjadi tensor
+
     with torch.no_grad():
         outputs = model(image_tensor)
+    
     predictions = outputs.argmax(dim=1)
     return predictions.item()
 
@@ -47,6 +47,7 @@ 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,
@@ -64,7 +65,6 @@ def get_nutritional_info(food):
                     nutrients_totals[nutrient_name] += nutrient_value
 
         average_nutrients = {name: total / item_count for name, total in nutrients_totals.items()}
-
         return average_nutrients
     else:
         return None