Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 from torchvision import transforms
 from transformers import (
     AutoImageProcessor, 
@@ -8,158 +8,178 @@ from transformers import (
     TrainingArguments,
     Trainer
 )
-from datasets import load_dataset, DatasetDict
+from datasets import load_dataset
 import numpy as np
 from huggingface_hub import HfApi
 import os
 from PIL import Image as PILImage
 
 # Configuration
-HF_DATASET = "Ultralytics/Brain-tumor"
 CUSTOM_MODEL_NAME = "GoGenix_MRI_Brain"
 BASE_MODEL = "Falconsai/nsfw_image_detection"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-def train_and_save_model():
-    """Train the model with explicit dataset format handling"""
+# Your custom dataset selection
+BRAIN_TUMOR_DATASETS = [
+    "PranomVignesh/MRI-Images-of-Brain-Tumor",  # Your first choice
+    "Hemg/Brain-Tumor-MRI-Dataset",           # Your second choice
+    "AntonXue/mcal-mri-brain-tumor",           # Your third choice
+]
+
+def find_working_dataset():
+    """Try your custom datasets and return the first one that works"""
+    working_datasets = []
     
-    try:
-        print("Loading Ultralytics/Brain-tumor dataset with explicit format...")
-        
-        # Try multiple loading methods to handle format detection issues
-        dataset = None
-        
-        # Method 1: Try loading with explicit imagefolder format for all splits
+    for dataset_name in BRAIN_TUMOR_DATASETS:
         try:
-            dataset = load_dataset(HF_DATASET, "imagefolder")
-            print("✅ Loaded with 'imagefolder' format")
-        except Exception as e1:
-            print(f"❌ Method 1 failed: {e1}")
+            print(f"Trying dataset: {dataset_name}")
+            dataset = load_dataset(dataset_name)
+            
+            # Basic validation
+            splits = list(dataset.keys())
+            print(f"Found splits: {splits}")
+            
+            # Check if dataset has content
+            if len(splits) == 0:
+                print(f"⚠️ {dataset_name} - No splits found")
+                continue
+                
+            first_split = splits[0]
+            if len(dataset[first_split]) == 0:
+                print(f"⚠️ {dataset_name} - Empty dataset")
+                continue
+            
+            # Check sample structure
+            sample = dataset[first_split][0]
+            sample_keys = list(sample.keys())
+            print(f"Sample keys: {sample_keys}")
             
-            # Method 2: Try loading without specific format
-            try:
-                dataset = load_dataset(HF_DATASET)
-                print("✅ Loaded without specific format")
-            except Exception as e2:
-                print(f"❌ Method 2 failed: {e2}")
+            if 'image' in sample_keys:
+                working_datasets.append({
+                    'name': dataset_name,
+                    'splits': splits,
+                    'sample_structure': sample_keys,
+                    'dataset': dataset
+                })
+                print(f"✅ {dataset_name} - VALID")
+            else:
+                print(f"⚠️ {dataset_name} - Missing 'image' key")
                 
-                # Method 3: Try loading with data_files specification
-                try:
-                    dataset = load_dataset(
-                        HF_DATASET,
-                        data_files={
-                            'train': ['**/train/**/*.jpg', '**/train/**/*.png', '**/train/**/*.jpeg'],
-                            'validation': ['**/valid/**/*.jpg', '**/valid/**/*.png', '**/valid/**/*.jpeg'],
-                            'test': ['**/test/**/*.jpg', '**/test/**/*.png', '**/test/**/*.jpeg']
-                        }
-                    )
-                    print("✅ Loaded with data_files specification")
-                except Exception as e3:
-                    print(f"❌ Method 3 failed: {e3}")
-                    return f"All loading methods failed:\n1. {e1}\n2. {e2}\n3. {e3}"
-        
-        if dataset is None:
-            return "❌ Could not load dataset with any method"
-        
-        print(f"Dataset splits available: {list(dataset.keys())}")
-        
-        # Check which splits we have and map them appropriately
-        if 'train' in dataset and 'validation' in dataset:
-            train_split = dataset['train']
-            test_split = dataset['validation']
-            print("Using 'train' and 'validation' splits")
-        elif 'valid' in dataset and 'test' in dataset:
-            train_split = dataset['valid']
-            test_split = dataset['test']
-            print("Using 'valid' and 'test' splits")
-        elif 'train' in dataset and 'test' in dataset:
-            train_split = dataset['train']
-            test_split = dataset['test']
-            print("Using 'train' and 'test' splits")
+        except Exception as e:
+            print(f"❌ {dataset_name} - Failed: {str(e)}")
+            continue
+    
+    return working_datasets
+
+def train_and_save_model():
+    """Train the model using your selected datasets"""
+    
+    try:
+        print("Searching for compatible brain tumor datasets...")
+        working_datasets = find_working_dataset()
+        
+        if not working_datasets:
+            return "❌ None of your selected datasets worked. Please check the dataset names or try different datasets."
+        
+        # Use the first working dataset
+        selected_dataset = working_datasets[0]
+        dataset_name = selected_dataset['name']
+        splits = selected_dataset['splits']
+        dataset_obj = selected_dataset['dataset']
+        
+        result_message = f"✅ Selected dataset: {dataset_name}\n"
+        result_message += f"Splits available: {splits}\n\n"
+        
+        print(f"Using dataset: {dataset_name}")
+        
+        # Determine which splits to use
+        train_split_key = None
+        test_split_key = None
+        
+        # Prioritize standard split names
+        if 'train' in splits:
+            train_split_key = 'train'
+        elif 'training' in splits:
+            train_split_key = 'training'
+        elif 'Train' in splits:
+            train_split_key = 'Train'
+        else:
+            train_split_key = splits[0]  # Use first available split
+        
+        if 'test' in splits:
+            test_split_key = 'test'
+        elif 'validation' in splits:
+            test_split_key = 'validation'
+        elif 'valid' in splits:
+            test_split_key = 'valid'
+        elif 'Test' in splits:
+            test_split_key = 'Test'
+        elif len(splits) > 1:
+            test_split_key = splits[1]  # Use second split
         else:
-            available_splits = list(dataset.keys())
-            return f"❌ Cannot determine train/test splits. Available splits: {available_splits}"
+            test_split_key = splits[0]  # Use same split for train/test (will split later)
         
-        print(f"Training samples: {len(train_split)}")
-        print(f"Test samples: {len(test_split)}")
+        train_split = dataset_obj[train_split_key]
+        test_split = dataset_obj[test_split_key]
         
-        # Analyze dataset structure
+        result_message += f"Using '{train_split_key}' split for training ({len(train_split)} samples)\n"
+        result_message += f"Using '{test_split_key}' split for testing ({len(test_split)} samples)\n\n"
+        
+        # Analyze dataset in detail
         if len(train_split) > 0:
             sample = train_split[0]
-            print(f"Sample keys: {list(sample.keys())}")
-            for key in sample.keys():
-                print(f"  {key}: {type(sample[key])}")
-        
-        # Determine if this is a classification or object detection dataset
-        # For Ultralytics datasets, check if it has object detection format
-        def detect_dataset_type(split):
-            if len(split) == 0:
-                return "empty"
+            result_message += f"Sample structure: {list(sample.keys())}\n"
             
-            sample = split[0]
-            if 'objects' in sample or 'bbox' in sample or 'labels' in sample and isinstance(sample.get('labels'), list):
-                return "object_detection"
-            elif 'label' in sample and isinstance(sample['label'], (int, float)):
-                return "classification"
-            elif 'image' in sample:
-                return "image_only"
-            else:
-                return "unknown"
-        
-        train_type = detect_dataset_type(train_split)
-        test_type = detect_dataset_type(test_split)
-        
-        print(f"Train dataset type: {train_type}")
-        print(f"Test dataset type: {test_type}")
-        
-        # Convert to classification format
-        def convert_to_classification(item):
-            """Convert various formats to classification format"""
-            image = item.get('image')
+            # Check image properties
+            if 'image' in sample:
+                img = sample['image']
+                result_message += f"Image type: {type(img)}, size: {getattr(img, 'size', 'N/A')}\n"
             
-            # Handle different label formats
-            if 'label' in item and isinstance(item['label'], (int, float)):
-                label = int(item['label'])
-            elif 'objects' in item or 'bbox' in item:
-                # Object detection format - convert to binary classification
-                # If there are objects/bboxes, it's tumor (1), else no tumor (0)
-                label = 1 if (item.get('objects') or item.get('bbox')) else 0
-            elif 'labels' in item and isinstance(item['labels'], list) and len(item['labels']) > 0:
-                label = 1  # Has labels = tumor
+            # Detect number of classes
+            if 'label' in sample:
+                unique_labels = set()
+                # Check first 100 samples for unique labels
+                for i in range(min(100, len(train_split))):
+                    unique_labels.add(train_split[i]['label'])
+                
+                num_classes = len(unique_labels)
+                result_message += f"Detected {num_classes} unique labels: {sorted(unique_labels)}\n"
+                
+                # Try to get class names
+                if hasattr(train_split.features.get('label', None), 'names'):
+                    class_names = train_split.features['label'].names
+                else:
+                    # Map numeric labels to meaningful names
+                    if num_classes == 2:
+                        class_names = ["no_tumor", "tumor"]
+                    elif num_classes == 3:
+                        class_names = ["glioma", "meningioma", "pituitary"]
+                    elif num_classes == 4:
+                        class_names = ["glioma", "meningioma", "no_tumor", "pituitary"]
+                    else:
+                        class_names = [f"class_{i}" for i in range(num_classes)]
+                
+                result_message += f"Using class names: {class_names}\n"
             else:
-                label = 0  # No labels = no tumor
-            
-            return {
-                'image': image,
-                'label': label
-            }
-        
-        print("Converting dataset to classification format...")
-        train_classification = train_split.map(convert_to_classification)
-        test_classification = test_split.map(convert_to_classification)
-        
-        # Count classes
-        tumor_count = sum(1 for item in train_classification if item['label'] == 1)
-        no_tumor_count = sum(1 for item in train_classification if item['label'] == 0)
-        
-        print(f"Tumor samples: {tumor_count}, No tumor samples: {no_tumor_count}")
-        
-        # Use binary classification
-        class_names = ["no_tumor", "tumor"]
-        num_classes = 2
-        
-        print(f"Using {num_classes} classes: {class_names}")
+                # Default to binary classification
+                num_classes = 2
+                class_names = ["no_tumor", "tumor"]
+                result_message += "No labels found, using binary classification\n"
+        else:
+            num_classes = 2
+            class_names = ["no_tumor", "tumor"]
+            result_message += "Empty dataset, using default binary classification\n"
         
-        # Define transforms
+        # Define transforms for MRI images
         transform = transforms.Compose([
             transforms.Resize((224, 224)),
-            transforms.Grayscale(num_output_channels=3),
+            transforms.Grayscale(num_output_channels=3),  # Ensure 3 channels
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
         
-        # Custom dataset class
-        class MRIDataset(torch.utils.data.Dataset):
+        # Custom dataset class with robust error handling
+        class MRIDataset(Dataset):
             def __init__(self, dataset, transform=None):
                 self.dataset = dataset
                 self.transform = transform
@@ -169,28 +189,42 @@ def train_and_save_model():
                 
             def __getitem__(self, idx):
                 item = self.dataset[idx]
-                image = item['image']
-                label = item['label']
                 
-                if self.transform:
-                    if not isinstance(image, PILImage.Image):
+                # Handle image
+                image = item.get('image')
+                if image is None:
+                    # Create placeholder image if none exists
+                    image = PILImage.new('RGB', (224, 224), color='gray')
+                elif not isinstance(image, PILImage.Image):
+                    try:
                         image = PILImage.fromarray(image)
+                    except:
+                        image = PILImage.new('RGB', (224, 224), color='gray')
+                
+                # Handle label
+                label = item.get('label', 0)
+                if isinstance(label, (list, tuple)) and len(label) > 0:
+                    label = label[0]  # Take first element if label is a list
+                label = int(label) if label is not None else 0
+                
+                # Apply transform
+                if self.transform:
                     image = self.transform(image)
                     
                 return image, label
         
         # Create datasets
-        train_dataset = MRIDataset(train_classification, transform=transform)
-        test_dataset = MRIDataset(test_classification, transform=transform)
+        train_dataset_obj = MRIDataset(train_split, transform=transform)
+        test_dataset_obj = MRIDataset(test_split, transform=transform)
         
-        # Load model
+        result_message += f"Final dataset - Train: {len(train_dataset_obj)}, Test: {len(test_dataset_obj)}\n\n"
+        
+        # Load base model
         print("Loading base model...")
         model = AutoModelForImageClassification.from_pretrained(
             BASE_MODEL,
             num_labels=num_classes,
-            ignore_mismatched_sizes=True,
-            id2label={0: "no_tumor", 1: "tumor"},
-            label2id={"no_tumor": 0, "tumor": 1}
+            ignore_mismatched_sizes=True
         )
         processor = AutoImageProcessor.from_pretrained(BASE_MODEL)
         model.to(DEVICE)
@@ -198,10 +232,10 @@ def train_and_save_model():
         # Training arguments
         training_args = TrainingArguments(
             output_dir="./results",
-            num_train_epochs=5,  # Reduced for testing
+            num_train_epochs=10,
             per_device_train_batch_size=8,
             per_device_eval_batch_size=8,
-            warmup_steps=100,
+            warmup_steps=500,
             weight_decay=0.01,
             logging_dir="./logs",
             logging_steps=10,
@@ -210,10 +244,9 @@ def train_and_save_model():
             load_best_model_at_end=True,
             push_to_hub=True,
             hub_model_id=CUSTOM_MODEL_NAME,
-            remove_unused_columns=False,
         )
         
-        # Metrics
+        # Metrics function
         def compute_metrics(eval_pred):
             predictions, labels = eval_pred
             predictions = np.argmax(predictions, axis=1)
@@ -224,79 +257,156 @@ def train_and_save_model():
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=train_dataset,
-            eval_dataset=test_dataset,
+            train_dataset=train_dataset_obj,
+            eval_dataset=test_dataset_obj,
             compute_metrics=compute_metrics,
         )
         
         # Start training
-        print("Starting training...")
+        result_message += "Starting training...\n"
         train_result = trainer.train()
         
         # Save model
         trainer.save_model(f"./{CUSTOM_MODEL_NAME}")
         processor.save_pretrained(f"./{CUSTOM_MODEL_NAME}")
         
-        # Push to hub
-        trainer.push_to_hub(commit_message="Train Brain Tumor classifier")
+        # Push to Hugging Face Hub
+        trainer.push_to_hub(commit_message=f"Trained on {dataset_name}")
         
-        # Results
+        # Training results
         train_accuracy = train_result.metrics.get('train_accuracy', 'N/A')
         eval_accuracy = train_result.metrics.get('eval_accuracy', 'N/A')
         
-        result_message = f"""
-        ✅ Training completed!
+        result_message += f"""
+        🎯 Training Completed Successfully!
         
+        Dataset: {dataset_name}
         Model: {CUSTOM_MODEL_NAME}
-        Dataset: {HF_DATASET}
         Classes: {class_names}
-        Training Accuracy: {train_accuracy}
-        Validation Accuracy: {eval_accuracy}
+        Training Accuracy: {train_accuracy or 'N/A'}
+        Validation Accuracy: {eval_accuracy or 'N/A'}
+        
+        Model has been saved and pushed to Hugging Face Hub.
+        You can now use the 'Classify MRI' tab to test the model.
         """
         
         return result_message
         
     except Exception as e:
         import traceback
-        return f"❌ Error: {str(e)}\n\n{traceback.format_exc()}"
+        error_details = traceback.format_exc()
+        
+        error_msg = f"""
+        ❌ Training Failed
+        
+        Error: {str(e)}
+        
+        Datasets tried:
+        {BRAIN_TUMOR_DATASETS}
+        
+        Please check:
+        1. Dataset names are correct
+        2. Internet connection
+        3. Dataset accessibility
+        
+        Error Details:
+        {error_details}
+        """
+        return error_msg
 
 def classify_mri(image):
-    """Classify MRI image"""
+    """Classify a new MRI image using the trained model"""
     try:
+        # Load your custom model
         model = AutoModelForImageClassification.from_pretrained(CUSTOM_MODEL_NAME)
         processor = AutoImageProcessor.from_pretrained(CUSTOM_MODEL_NAME)
         
         model.to(DEVICE)
         model.eval()
         
+        # Preprocess image
         inputs = processor(image, return_tensors="pt").to(DEVICE)
         
+        # Predict
         with torch.no_grad():
             outputs = model(**inputs)
             predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
         
-        class_names = ["No Tumor", "Tumor Detected"]
-        results = {class_names[i]: float(predictions[0][i]) for i in range(2)}
+        # Try to detect number of classes
+        num_classes = predictions.shape[1]
+        
+        # Default class names based on number of classes
+        if num_classes == 2:
+            class_names = ["No Tumor", "Tumor Detected"]
+        elif num_classes == 3:
+            class_names = ["Glioma", "Meningioma", "Pituitary Tumor"]
+        elif num_classes == 4:
+            class_names = ["Glioma", "Meningioma", "No Tumor", "Pituitary Tumor"]
+        else:
+            class_names = [f"Class {i}" for i in range(num_classes)]
+        
+        results = {class_names[i]: float(predictions[0][i]) for i in range(num_classes)}
         
         return results
     
     except Exception as e:
-        return f"⚠️ Error: {str(e)}"
+        return f"⚠️ Model not trained yet or unavailable. Error: {str(e)}"
 
-# Simple Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Brain Tumor Classification")
+# Gradio Interface
+with gr.Blocks(title="GoGenix MRI Brain Tumor Classifier") as demo:
+    gr.Markdown("# 🧠 GoGenix MRI Brain Tumor Classifier")
+    gr.Markdown("**Using Your Selected Datasets**")
+    
+    with gr.Tab("🚀 Train Model"):
+        gr.Markdown("### Train with Your Custom Datasets")
+        gr.Markdown("Will try these datasets in order:")
+        for i, dataset in enumerate(BRAIN_TUMOR_DATASETS, 1):
+            gr.Markdown(f"{i}. `{dataset}`")
+        
+        train_btn = gr.Button("Start Training", variant="primary", size="lg")
+        output_text = gr.Textbox(
+            label="Training Status", 
+            lines=20,
+            placeholder="Click 'Start Training' to begin..."
+        )
+        
+        train_btn.click(
+            fn=train_and_save_model,
+            outputs=output_text
+        )
     
-    with gr.Tab("Train"):
-        train_btn = gr.Button("Train Model")
-        output = gr.Textbox(lines=10)
-        train_btn.click(train_and_save_model, outputs=output)
+    with gr.Tab("🔍 Classify MRI"):
+        gr.Markdown("### Upload MRI Image for Classification")
+        gr.Markdown("**Note**: Requires successful training first")
+        
+        image_input = gr.Image(
+            type="pil", 
+            label="Brain MRI Scan",
+            height=300
+        )
+        classify_btn = gr.Button("Classify", variant="secondary")
+        result = gr.Label(
+            label="Brain Tumor Classification Results",
+            num_top_classes=4
+        )
+        
+        classify_btn.click(
+            fn=classify_mri,
+            inputs=image_input,
+            outputs=result
+        )
     
-    with gr.Tab("Classify"):
-        image = gr.Image(type="pil")
-        classify_btn = gr.Button("Classify")
-        result = gr.Label()
-        classify_btn.click(classify_mri, inputs=image, outputs=result)
+    with gr.Tab("📊 Your Datasets"):
+        gr.Markdown("### Your Selected Brain Tumor Datasets")
+        gr.Markdown("""
+        **Currently Using:**
+        
+        1. **PranomVignesh/MRI-Images-of-Brain-Tumor** - Primary choice
+        2. **Hemg/Brain-Tumor-MRI-Dataset** - Secondary choice  
+        3. **AntonXue/mcal-mri-brain-tumor** - Tertiary choice
+        
+        The system will try these in order and use the first one that works.
+        """)
 
 if __name__ == "__main__":
     demo.launch()