update app.py with nutritional dataset

app.py

@@ -3,12 +3,80 @@ import model_builder as mb
 from torchvision import transforms
 import torch 
 
-device = torch.device("cpu")
+# Comprehensive nutrition data per 165g serving
+NUTRITION_DATA = {
+    'Fresh Apple': {
+        'macronutrients': {
+            'calories': 99.2,
+            'protein': 0.8,
+            'carbs': 23.3,
+            'fats': 0.3,
+            'water': 140.2,
+            'fiber': 1.5
+        },
+        'micronutrients': {
+            'vitamin_c': 96.7,
+            'thiamin': 0.1,
+            'niacin': 0.4,
+            'vitamin_b6': 0.2
+        },
+        'macrominerals': {
+            'magnesium': 22.1,
+            'phosphorus': 8.9,
+            'potassium': 226.0,
+            'calcium': 20.6
+        }
+    },
+    'Fresh Banana': {
+        'macronutrients': {
+            'calories': 147.0,
+            'protein': 1.8,
+            'carbs': 38.0,
+            'fats': 0.5,
+            'water': 132.0,
+            'fiber': 3.5
+        },
+        'micronutrients': {
+            'vitamin_c': 14.7,
+            'thiamin': 0.4,
+            'niacin': 1.2,
+            'vitamin_b6': 0.5
+        },
+        'macrominerals': {
+            'magnesium': 41.3,
+            'phosphorus': 33.0,
+            'potassium': 537.0,
+            'calcium': 8.3
+        }
+    },
+    'Fresh Orange': {
+        'macronutrients': {
+            'calories': 82.0,
+            'protein': 1.6,
+            'carbs': 21.0,
+            'fats': 0.2,
+            'water': 146.0,
+            'fiber': 4.0
+        },
+        'micronutrients': {
+            'vitamin_c': 82.7,
+            'thiamin': 0.2,
+            'niacin': 0.5,
+            'vitamin_b6': 0.1
+        },
+        'macrominerals': {
+            'magnesium': 18.2,
+            'phosphorus': 28.1,
+            'potassium': 237.6,
+            'calcium': 74.3
+        }
+    }
+}
 
+device = torch.device("cpu")
 normalize = transforms.Normalize(
         mean=[0.485, 0.456, 0.406], 
         std=[0.229, 0.224, 0.225])
-
 manual_transform = transforms.Compose([
     transforms.ToPILImage(), 
     transforms.Resize(size=(224, 224)), 
@@ -16,23 +84,6 @@ manual_transform = transforms.Compose([
     normalize
 ])
 
-# class_names = ['Fresh Banana',
-#   'Fresh Lemon',
-#   'Fresh Lulo',
-#   'Fresh Mango',
-#   'Fresh Orange',
-#   'Fresh Strawberry',
-#   'Fresh Tamarillo',
-#   'Fresh Tomato',
-#   'Spoiled Banana',
-#   'Spoiled Lemon',
-#   'Spoiled Lulo',
-#   'Spoiled Mango',
-#   'Spoiled Orange',
-#   'Spoiled Strawberry',
-#   'Spoiled Tamarillo',
-#   'Spoiled Tomato']
-
 class_names = ['Fresh Apple',
  'Fresh Banana',
  'Fresh Orange',
@@ -43,13 +94,58 @@ class_names = ['Fresh Apple',
 model = mb.create_model_baseline_effnetb0(out_feats=len(class_names), device=device)
 model.load_state_dict(torch.load(f="models/effnetb0_freshvisionv0_10_epochs.pt", map_location="cpu"))
 
+def format_nutrition(fruit_name):
+    """Format comprehensive nutrition information for display"""
+    if fruit_name not in NUTRITION_DATA:
+        return ""
+    
+    nutrition = NUTRITION_DATA[fruit_name]
+    macro = nutrition['macronutrients']
+    micro = nutrition['micronutrients']
+    minerals = nutrition['macrominerals']
+    
+    return f"""
+    Nutritional Information (per 165g serving):
+    
+    Macronutrients:
+    • Calories: {macro['calories']} kcal
+    • Protein: {macro['protein']} g
+    • Carbs: {macro['carbs']} g
+    • Fats: {macro['fats']} g
+    • Water: {macro['water']} ml
+    • Fiber: {macro['fiber']} g
+    
+    Micronutrients:
+    • Vitamin C: {micro['vitamin_c']} mg
+    • Thiamin: {micro['thiamin']} mg
+    • Niacin: {micro['niacin']} mg
+    • Vitamin B6: {micro['vitamin_b6']} mg
+    
+    Macrominerals:
+    • Magnesium: {minerals['magnesium']} mg
+    • Phosphorus: {minerals['phosphorus']} mg
+    • Potassium: {minerals['potassium']} mg
+    • Calcium: {minerals['calcium']} mg
+    """
+
 def pred(img):
     model.eval()
     transformed = manual_transform(img).to(device)
     with torch.inference_mode():
         logits = model(transformed.unsqueeze(dim=0))
         pred = torch.softmax(logits, dim=-1)
-    return f"prediction: {class_names[pred.argmax(dim=-1).item()]} | confidence: {pred.max():.3f}"
+    
+    predicted_class = class_names[pred.argmax(dim=-1).item()]
+    confidence = pred.max().item()
+    
+    result = f"Prediction: {predicted_class} | Confidence: {confidence:.3f}"
+    
+    # Add nutrition information if it's a fresh fruit
+    if predicted_class.startswith('Fresh'):
+        nutrition_info = format_nutrition(predicted_class)
+        result += f"\n{nutrition_info}"
+    
+    return result
 
 demo = gr.Blocks()
 
@@ -61,11 +157,16 @@ with demo:
         This model has been trained on [kaggle datasets](https://www.kaggle.com/datasets/sriramr/fruits-fresh-and-rotten-for-classification) using NVIDIA T4 GPU._
         
         ## Model capabilities:
-        - Classify freshness from fruits image (apple, orange, and banana) with two labels: _Fresh_ and _Rotten/spoiled_
+        - Classify freshness from fruits image (apple, orange, and banana) with two labels: *Fresh* and *Rotten/spoiled*
+        - Provides comprehensive nutritional information for fresh fruits including:
+          * Macronutrients (calories, protein, carbs, fats, water, fiber)
+          * Micronutrients (vitamins C, B6, thiamin, niacin)
+          * Macrominerals (magnesium, phosphorus, potassium, calcium)
+        
         ## Model drawbacks: 
         - Sometimes perform false prediction on some fruits condition, this is due to low variability on the image datasets
         - Can't perform accurate prediction on multiple objects/combined condition (e.g. two bananas with different freshness condition)
-        - This models can't identify non-fruits objects , since it's only trained with fruits dataset. 
+        - This models can't identify non-fruits objects, since it's only trained with fruits dataset
                 
         ## **How to get the best result with this model:** 
         1. The image should only contain fruits that the model can recognize (apple, orange, and banana)
@@ -73,8 +174,13 @@ with demo:
         3. Ensure the object is captured with sufficient light so that the surface of the fruits is exposed properly
                 
         get the [source code](https://github.com/devdezzies/freshvision) on my github
-""")
-    gr.Interface(pred, gr.Image(), outputs="text")
+    """)
+    gr.Interface(
+        fn=pred,
+        inputs=gr.Image(),
+        outputs=gr.Textbox(label="Prediction Results", lines=15),
+        title="FreshVision Fruit Classifier"
+    )
 
 if __name__ == "__main__":
     demo.launch()