Update app.py

app.py

@@ -10,36 +10,36 @@ from torchvision import transforms, models
 from torch import nn, optim
 from torch.utils.data import Dataset, DataLoader
 import time
-import os
+import random
 
-# Set device
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Set device - use CPU for Hugging Face Spaces
+device = torch.device('cpu')
 print(f"Using device: {device}")
 
-# Step 1: Import the dataset
-print("🚀 Loading MM-Food-100K dataset...")
-dataset = load_dataset("Codatta/MM-Food-100K", split="train", streaming=True)
-
-# Take a sample to explore (first 1000 images for demonstration)
-print("📦 Taking sample of 1000 images for exploration...")
-sample_data = list(dataset.take(1000))
+# Step 1: Import the CORRECT dataset with actual images
+print("🚀 Loading Food-101 dataset...")
+# Load a smaller sample for Hugging Face Spaces constraints
+dataset = load_dataset("ethz/food101", split="train[:5000]")  # Take first 5000 images for demo
 
 # Step 2: Check the rows and structure
 print("\n📋 Dataset structure sample:")
-print(f"Total samples in full dataset: 100,000")
-print(f"Sample size we're using: {len(sample_data)}")
+print(f"Total samples in our sample: {len(dataset)}")
 print("\nFirst row example:")
-print(sample_data[0])
+print(dataset[0])
 
-# Extract images and labels for exploration
-images = [item['image'] for item in sample_data]
-labels = [item['label'] for item in sample_data]  # Assuming there's a 'label' field
+# Extract images and labels
+images = [item['image'] for item in dataset]
+labels = [item['label'] for item in dataset]
 
 print(f"\n🔍 Unique classes in sample: {len(set(labels))}")
 print(f"📊 Class distribution:")
 class_counts = pd.Series(labels).value_counts()
 print(class_counts.head(10))
 
+# Get class names mapping
+id2label = dataset.features['label'].names
+num_classes = len(id2label)
+
 # Step 3: Visualize charts with matplotlib
 def create_visualizations():
     """Create exploratory visualizations"""
@@ -52,20 +52,9 @@ def create_visualizations():
     ax1.set_ylabel('Count')
     ax1.tick_params(axis='x', rotation=45)
     
-    # Image size distribution (if available)
-    img_sizes = [(img.size[0], img.size[1]) for img in images[:100]]
-    widths = [size[0] for size in img_sizes]
-    heights = [size[1] for size in img_sizes]
-    
-    ax2.scatter(widths, heights, alpha=0.6, color='green')
-    ax2.set_title('Image Dimensions Distribution')
-    ax2.set_xlabel('Width (pixels)')
-    ax2.set_ylabel('Height (pixels)')
-    ax2.grid(True, alpha=0.3)
-    
     # Sample images grid
-    ax3.set_title('Sample Food Images')
-    ax3.axis('off')
+    ax2.set_title('Sample Food Images')
+    ax2.axis('off')
     # Create a grid of sample images
     grid_img = Image.new('RGB', (400, 400), (255, 255, 255))
     for i in range(min(4, len(images))):
@@ -73,11 +62,22 @@ def create_visualizations():
         x = (i % 2) * 200
         y = (i // 2) * 200
         grid_img.paste(img, (x, y))
-    ax3.imshow(grid_img)
+    ax2.imshow(grid_img)
     
     # Class frequency pie chart
-    ax4.pie(class_counts.head(5), labels=class_counts.head(5).index, autopct='%1.1f%%', startangle=90)
-    ax4.set_title('Top 5 Classes Percentage')
+    ax3.pie(class_counts.head(5), labels=class_counts.head(5).index, autopct='%1.1f%%', startangle=90)
+    ax3.set_title('Top 5 Classes Percentage')
+    
+    # Image size distribution
+    img_sizes = [(img.size[0], img.size[1]) for img in images[:100]]
+    widths = [size[0] for size in img_sizes]
+    heights = [size[1] for size in img_sizes]
+    
+    ax4.scatter(widths, heights, alpha=0.6, color='green')
+    ax4.set_title('Image Dimensions Distribution')
+    ax4.set_xlabel('Width (pixels)')
+    ax4.set_ylabel('Height (pixels)')
+    ax4.grid(True, alpha=0.3)
     
     plt.tight_layout()
     return fig
@@ -89,21 +89,29 @@ print("✅ Visualizations created successfully!")
 
 # Step 4: Train/Test Split
 print("\n✂️ Creating train/test split...")
-train_images, test_images, train_labels, test_labels = train_test_split(
-    images, labels, test_size=0.2, random_state=42, stratify=labels if len(set(labels)) < len(labels) else None
-)
+train_size = min(4000, len(images) - 1000)  # Keep it small for Spaces
+test_size = min(1000, len(images) - train_size)
+
+train_images = images[:train_size]
+train_labels = labels[:train_size]
+test_images = images[train_size:train_size+test_size]
+test_labels = labels[train_size:train_size+test_size]
 
 print(f"Training samples: {len(train_images)}")
 print(f"Testing samples: {len(test_images)}")
-print(f"Number of classes: {len(set(labels))}")
+print(f"Number of classes: {num_classes}")
 
 # Step 5: Set up EfficientNet-B0 model
 class FoodClassifier(nn.Module):
-    def __init__(self, num_classes=50):  # Adjust based on actual classes
+    def __init__(self, num_classes=101):
         super(FoodClassifier, self).__init__()
         # Load pre-trained EfficientNet-B0
         self.effnet = models.efficientnet_b0(pretrained=True)
         
+        # Freeze most layers for demo training
+        for param in self.effnet.parameters():
+            param.requires_grad = False
+        
         # Replace the classifier head
         num_features = self.effnet.classifier[1].in_features
         self.effnet.classifier = nn.Sequential(
@@ -118,8 +126,8 @@ class FoodClassifier(nn.Module):
 train_transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.RandomHorizontalFlip(),
-    transforms.RandomRotation(10),
-    transforms.ColorJitter(brightness=0.2, contrast=0.2),
+    transforms.RandomRotation(5),
+    transforms.ColorJitter(brightness=0.1, contrast=0.1),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
@@ -136,7 +144,6 @@ class FoodDataset(Dataset):
         self.images = images
         self.labels = labels
         self.transform = transform
-        self.label_to_idx = {label: idx for idx, label in enumerate(set(labels))}
     
     def __len__(self):
         return len(self.images)
@@ -144,41 +151,44 @@ class FoodDataset(Dataset):
     def __getitem__(self, idx):
         img = self.images[idx]
         label = self.labels[idx]
-        label_idx = self.label_to_idx[label]
+        
+        # Convert to RGB if needed
+        if img.mode != 'RGB':
+            img = img.convert('RGB')
         
         if self.transform:
             img = self.transform(img)
         
-        return img, label_idx
+        return img, label
 
 # Create datasets
 train_dataset = FoodDataset(train_images, train_labels, transform=train_transform)
 test_dataset = FoodDataset(test_images, test_labels, transform=test_transform)
 
-# Create data loaders
-train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
-test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
+# Create data loaders with small batch sizes for Spaces
+train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)
 
 # Initialize model
-num_classes = len(set(labels))
 model = FoodClassifier(num_classes=num_classes).to(device)
 print(f"Intialized EfficientNet-B0 model with {num_classes} output classes")
 
 # Training function (simplified for demo)
 def train_model():
     """Simple training function for demonstration"""
-    print("\n🏋️ Starting model training...")
+    print("\n🏋️ Starting model training (demo mode)...")
     criterion = nn.CrossEntropyLoss()
-    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    optimizer = optim.Adam(model.effnet.classifier.parameters(), lr=0.001)  # Only train classifier
     
     model.train()
     total_loss = 0
     correct = 0
     total = 0
     
-    # Train on a small subset for demo purposes
+    # Train on a very small subset for demo
+    max_batches = 20  # Very limited for Spaces
     for batch_idx, (images, labels) in enumerate(train_loader):
-        if batch_idx >= 10:  # Only train on first 10 batches for demo
+        if batch_idx >= max_batches:
             break
             
         images, labels = images.to(device), labels.to(device)
@@ -194,11 +204,11 @@ def train_model():
         total += labels.size(0)
         correct += (predicted == labels).sum().item()
         
-        if batch_idx % 2 == 0:
-            print(f"Batch {batch_idx}/10 - Loss: {loss.item():.4f}")
+        if batch_idx % 5 == 0:
+            print(f"Batch {batch_idx}/{max_batches} - Loss: {loss.item():.4f}")
     
-    accuracy = 100 * correct / total
-    avg_loss = total_loss / min(10, len(train_loader))
+    accuracy = 100 * correct / total if total > 0 else 0
+    avg_loss = total_loss / min(max_batches, len(train_loader))
     print(f"✅ Training completed! Average Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%")
     return avg_loss, accuracy
 
@@ -206,6 +216,10 @@ def train_model():
 def predict_food(image):
     """Predict food class from uploaded image"""
     try:
+        # Convert to RGB if needed
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
         # Preprocess image
         img_tensor = test_transform(image).unsqueeze(0).to(device)
         
@@ -221,7 +235,8 @@ def predict_food(image):
         # Create results string
         results = []
         for i in range(top5_prob.size(0)):
-            class_name = list(train_dataset.label_to_idx.keys())[list(train_dataset.label_to_idx.values()).index(top5_catid[i].item())]
+            class_idx = top5_catid[i].item()
+            class_name = id2label[class_idx] if class_idx < len(id2label) else f"Class {class_idx}"
             probability = top5_prob[i].item() * 100
             results.append(f"{class_name}: {probability:.2f}%")
         
@@ -266,10 +281,9 @@ with gr.Blocks(title="Food Classifier") as demo:
     with gr.Tabs():
         with gr.TabItem("📊 Dataset Explorer"):
             gr.Markdown("## Dataset Exploration")
-            gr.Markdown(f"- **Dataset**: Codatta/MM-Food-100K")
-            gr.Markdown(f"- **Total Images**: 100,000")
-            gr.Markdown(f"- **Sample Size Used**: 1,000")
-            gr.Markdown(f"- **Unique Classes**: {len(set(labels))}")
+            gr.Markdown(f"- **Dataset**: Food-101 (sample)") 
+            gr.Markdown(f"- **Total Images**: 101,000 (using 5,000 sample)")
+            gr.Markdown(f"- **Unique Classes**: {num_classes}")
             gr.Plot(vis_fig)
         
         with gr.TabItem("🎯 Classifier"):
@@ -289,10 +303,13 @@ with gr.Blocks(title="Food Classifier") as demo:
     )
     
     gr.Markdown("### Training Summary")
-    gr.Markdown(f"- **Model**: EfficientNet-B0")
+    gr.Markdown(f"- **Model**: EfficientNet-B0 (transfer learning)")
     gr.Markdown(f"- **Training Accuracy**: {train_acc:.2f}%")
     gr.Markdown(f"- **Training Loss**: {train_loss:.4f}")
     gr.Markdown("- **Note**: This is a demo training on a small subset. For full training, use Google Colab/Kaggle.")
+    
+    gr.Markdown("### Dataset Information")
+    gr.Markdown("- This dataset consists of 101 food categories, with 101'000 images. For each class, 250 manually reviewed test images are provided as well as 750 training images") [[1]]
 
 # Launch the app
 if __name__ == "__main__":