Update app.py

app.py

@@ -8,119 +8,217 @@ from transformers import (
     TrainingArguments,
     Trainer
 )
-from datasets import load_dataset, Image
+from datasets import load_dataset, Dataset, Image
 import numpy as np
-from huggingface_hub import HfApi, create_repo
+from huggingface_hub import HfApi
 import os
+import json
+from PIL import Image as PILImage
 
 # Configuration
-HF_DATASET = "Ultralytics/Brain-tumor"  # Your dataset repo
-CUSTOM_MODEL_NAME = "GoGenix_MRI_Brain"              # Your custom model name
+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")
 
-class MRIDataset(torch.utils.data.Dataset):
-    def __init__(self, dataset, transform=None):
-        self.dataset = dataset
-        self.transform = transform
+def train_and_save_model():
+    """Train the model using YOLO format dataset"""
+    
+    try:
+        print("Loading Ultralytics/Brain-tumor dataset (YOLO format)...")
+        
+        # Load the dataset
+        dataset = load_dataset(HF_DATASET)
         
-    def __len__(self):
-        return len(self.dataset)
+        print(f"Dataset splits available: {list(dataset.keys())}")
         
-    def __getitem__(self, idx):
-        item = self.dataset[idx]
-        image = item['image']
-        label = item['label']
+        # Check dataset structure
+        if 'valid' not in dataset or 'test' not in dataset:
+            return "❌ Error: Dataset must contain 'valid' and 'test' splits"
         
-        if self.transform:
-            image = self.transform(image)
+        train_split = dataset['valid']
+        test_split = dataset['test']
+        
+        print("Analyzing YOLO dataset structure...")
+        
+        # For YOLO datasets, we need to check if images and labels are separate
+        # Let's examine the structure
+        if len(train_split) > 0:
+            sample = train_split[0]
+            print(f"Sample keys: {list(sample.keys())}")
             
-        return image, label
-
-def train_and_save_model():
-    """Train the model and save as GoGenix_MRI_Brain"""
-    
-    # Load dataset from Hugging Face Hub
-    print("Loading dataset from Hugging Face Hub...")
-    dataset = load_dataset(HF_DATASET)
-    
-    # Get class names from dataset
-    class_names = dataset['train'].features['label'].names
-    print(f"Classes detected: {class_names}")
-    
-    # Define transforms for MRI images
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.Grayscale(num_output_channels=3),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-    
-    # Create PyTorch datasets
-    train_dataset = MRIDataset(dataset['train'], transform=transform)
-    test_dataset = MRIDataset(dataset['test'], transform=transform)
-    
-    # Load base model
-    print("Loading base model...")
-    model = AutoModelForImageClassification.from_pretrained(
-        BASE_MODEL,
-        num_labels=len(class_names),
-        ignore_mismatched_sizes=True
-    )
-    processor = AutoImageProcessor.from_pretrained(BASE_MODEL)
-    
-    model.to(DEVICE)
-    
-    # Training arguments
-    training_args = TrainingArguments(
-        output_dir="./results",
-        num_train_epochs=10,
-        per_device_train_batch_size=8,
-        per_device_eval_batch_size=8,
-        warmup_steps=500,
-        weight_decay=0.01,
-        logging_dir="./logs",
-        logging_steps=10,
-        evaluation_strategy="epoch",
-        save_strategy="epoch",
-        load_best_model_at_end=True,
-        push_to_hub=True,
-        hub_model_id=CUSTOM_MODEL_NAME,
-    )
-    
-    # Custom compute_metrics function
-    def compute_metrics(eval_pred):
-        predictions, labels = eval_pred
-        predictions = np.argmax(predictions, axis=1)
-        return {"accuracy": (predictions == labels).mean()}
-    
-    # Create trainer
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=test_dataset,
-        compute_metrics=compute_metrics,
-    )
-    
-    # Start training
-    print("Starting training...")
-    train_result = trainer.train()
-    
-    # Save metrics
-    trainer.log_metrics("train", train_result.metrics)
-    trainer.save_metrics("train", train_result.metrics)
-    trainer.save_state()
-    
-    # Save model locally
-    trainer.save_model(f"./{CUSTOM_MODEL_NAME}")
-    processor.save_pretrained(f"./{CUSTOM_MODEL_NAME}")
-    
-    # Push to Hugging Face Hub
-    print("Pushing model to Hugging Face Hub...")
-    trainer.push_to_hub()
-    
-    return f"Training completed! Model saved as: {CUSTOM_MODEL_NAME}"
+            # Check if it's YOLO format (has image path and labels path)
+            if 'image' in sample:
+                print(f"Image type: {type(sample['image'])}")
+            if 'label' in sample:
+                print(f"Label type: {type(sample['label'])}")
+                if isinstance(sample['label'], list) and len(sample['label']) > 0:
+                    print(f"First label sample: {sample['label'][0]}")
+        
+        # Since Ultralytics datasets are typically for object detection,
+        # we'll convert them to classification by checking if tumor is present
+        def yolo_to_classification(item):
+            """Convert YOLO object detection labels to classification labels"""
+            image = item['image']
+            labels = item.get('label', [])
+            
+            # For binary classification: 0 = no tumor, 1 = tumor present
+            # If there are any labels (bounding boxes), it means tumor is present
+            has_tumor = 1 if labels and len(labels) > 0 else 0
+            
+            return {
+                'image': image,
+                'label': has_tumor
+            }
+        
+        # Apply conversion
+        print("Converting YOLO labels to classification...")
+        train_classification = train_split.map(yolo_to_classification)
+        test_classification = test_split.map(yolo_to_classification)
+        
+        # Count tumor vs no_tumor
+        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"Training set - Tumors: {tumor_count}, No tumors: {no_tumor_count}")
+        
+        # Define class names for binary classification
+        class_names = ["no_tumor", "tumor"]
+        num_classes = 2
+        
+        print(f"Using binary classification: {class_names}")
+        
+        # Define transforms for MRI images
+        transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.Grayscale(num_output_channels=3),
+            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):
+            def __init__(self, dataset, transform=None):
+                self.dataset = dataset
+                self.transform = transform
+                
+            def __len__(self):
+                return len(self.dataset)
+                
+            def __getitem__(self, idx):
+                item = self.dataset[idx]
+                image = item['image']
+                label = item['label']
+                
+                if self.transform:
+                    # Ensure image is PIL Image
+                    if not isinstance(image, PILImage.Image):
+                        image = PILImage.fromarray(image)
+                    image = self.transform(image)
+                    
+                return image, label
+        
+        # Create dataset objects
+        train_dataset = MRIDataset(train_classification, transform=transform)
+        test_dataset = MRIDataset(test_classification, transform=transform)
+        
+        print(f"Training samples: {len(train_dataset)}")
+        print(f"Test samples: {len(test_dataset)}")
+        
+        # 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}
+        )
+        processor = AutoImageProcessor.from_pretrained(BASE_MODEL)
+        model.to(DEVICE)
+        
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir="./results",
+            num_train_epochs=10,
+            per_device_train_batch_size=8,
+            per_device_eval_batch_size=8,
+            warmup_steps=500,
+            weight_decay=0.01,
+            logging_dir="./logs",
+            logging_steps=10,
+            evaluation_strategy="epoch",
+            save_strategy="epoch",
+            load_best_model_at_end=True,
+            push_to_hub=True,
+            hub_model_id=CUSTOM_MODEL_NAME,
+            remove_unused_columns=False,
+        )
+        
+        # Metrics function
+        def compute_metrics(eval_pred):
+            predictions, labels = eval_pred
+            predictions = np.argmax(predictions, axis=1)
+            accuracy = (predictions == labels).mean()
+            return {"accuracy": accuracy}
+        
+        # Create trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=test_dataset,
+            compute_metrics=compute_metrics,
+        )
+        
+        # Start training
+        print("Starting training...")
+        train_result = trainer.train()
+        
+        # Save model locally
+        trainer.save_model(f"./{CUSTOM_MODEL_NAME}")
+        processor.save_pretrained(f"./{CUSTOM_MODEL_NAME}")
+        
+        # Push to Hugging Face Hub
+        print("Pushing model to Hugging Face Hub...")
+        trainer.push_to_hub(commit_message="Train Brain Tumor classifier (YOLO to Classification)")
+        
+        # Display 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 successfully!
+        
+        Model: {CUSTOM_MODEL_NAME}
+        Dataset: {HF_DATASET} (YOLO format)
+        Task: Binary Classification (Tumor Detection)
+        Classes: {', '.join(class_names)}
+        Training Samples: {len(train_dataset)}
+        Test Samples: {len(test_dataset)}
+        Training Accuracy: {train_accuracy}
+        Validation Accuracy: {eval_accuracy}
+        Tumor/No-Tumor Ratio: {tumor_count}/{no_tumor_count}
+        
+        Model has been saved and pushed to Hugging Face Hub.
+        """
+        
+        return result_message
+        
+    except Exception as e:
+        import traceback
+        error_details = traceback.format_exc()
+        
+        error_message = f"""
+        ❌ Error during training:
+        
+        Error Type: {type(e).__name__}
+        Error Message: {str(e)}
+        
+        Detailed Traceback:
+        {error_details}
+        """
+        return error_message
 
 def classify_mri(image):
     """Classify a new MRI image using the trained model"""
@@ -140,25 +238,47 @@ def classify_mri(image):
             outputs = model(**inputs)
             predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
         
-        # Map to class names
-        class_names = ["glioma", "meningioma", "no_tumor", "pituitary"]
-        return {class_names[i]: float(predictions[0][i]) for i in range(len(class_names))}
+        # Binary classification results
+        class_names = ["No Tumor", "Tumor Detected"]
+        results = {
+            class_names[0]: float(predictions[0][0]),  # No tumor probability
+            class_names[1]: float(predictions[0][1])   # Tumor probability
+        }
+        
+        # Add diagnostic information
+        tumor_prob = float(predictions[0][1])
+        if tumor_prob > 0.7:
+            diagnosis = "🟢 Likely no tumor"
+        elif tumor_prob > 0.3:
+            diagnosis = "🟡 Uncertain - consult specialist"
+        else:
+            diagnosis = "🔴 Possible tumor detected"
+        
+        return {
+            "classification": results,
+            "diagnosis": diagnosis,
+            "tumor_probability": tumor_prob
+        }
     
     except Exception as e:
-        return f"Error: Model not trained yet or unavailable. Please train first."
+        return f"⚠️ Model not trained yet or unavailable. Error: {str(e)}"
 
 # Gradio Interface
 with gr.Blocks(title="GoGenix MRI Brain Tumor Classifier") as demo:
     gr.Markdown("# 🧠 GoGenix MRI Brain Tumor Classifier")
-    gr.Markdown(f"Training **Falconsai/nsfw_image_detection** on Brain Tumor MRI dataset")
+    gr.Markdown(f"**Dataset**: {HF_DATASET} (YOLO Format) | **Base Model**: {BASE_MODEL}")
     
     with gr.Tab("🚀 Train Model"):
         gr.Markdown("### Train GoGenix_MRI_Brain Model")
-        gr.Markdown(f"Dataset: `{HF_DATASET}`")
-        gr.Markdown(f"Target Model: `{CUSTOM_MODEL_NAME}`")
+        gr.Markdown(f"Using YOLO format dataset: `{HF_DATASET}`")
+        gr.Markdown("**Note**: Converting object detection labels to binary classification")
         
-        train_btn = gr.Button("Start Training", variant="primary")
-        output_text = gr.Textbox(label="Training Status", lines=5)
+        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,
@@ -166,26 +286,54 @@ with gr.Blocks(title="GoGenix MRI Brain Tumor Classifier") as demo:
         )
     
     with gr.Tab("🔍 Classify MRI"):
-        gr.Markdown("### Upload MRI Image for Classification")
-        image_input = gr.Image(type="pil", label="MRI Scan")
-        classify_btn = gr.Button("Classify", variant="secondary")
-        result = gr.Label(label="Tumor Classification Results")
+        gr.Markdown("### Upload MRI Image for Tumor Detection")
+        gr.Markdown("**Binary Classification**: Tumor vs No Tumor")
+        
+        image_input = gr.Image(
+            type="pil", 
+            label="Brain MRI Scan",
+            height=300
+        )
+        classify_btn = gr.Button("Analyze Scan", variant="secondary")
+        
+        with gr.Row():
+            result_label = gr.Label(
+                label="Classification Results",
+                num_top_classes=2
+            )
+            diagnosis_text = gr.Textbox(
+                label="Diagnostic Suggestion",
+                interactive=False
+            )
+        
+        def process_classification(image):
+            result = classify_mri(image)
+            if isinstance(result, dict) and 'classification' in result:
+                return result['classification'], result.get('diagnosis', '')
+            else:
+                return {"Error": 1.0}, result
         
         classify_btn.click(
-            fn=classify_mri,
+            fn=process_classification,
             inputs=image_input,
-            outputs=result
+            outputs=[result_label, diagnosis_text]
         )
     
-    with gr.Tab("📊 Model Info"):
-        gr.Markdown("### Model Information")
+    with gr.Tab("📊 Dataset Info"):
+        gr.Markdown("### YOLO Dataset Information")
         gr.Markdown(f"""
-        - **Base Model**: {BASE_MODEL}
-        - **Custom Model**: {CUSTOM_MODEL_NAME}
-        - **Dataset**: {HF_DATASET}
-        - **Classes**: Glioma, Meningioma, No Tumor, Pituitary Tumor
-        - **Device**: {DEVICE}
+        **Dataset**: {HF_DATASET}
+        **Format**: YOLO (You Only Look Once) Object Detection
+        **Original Structure**:
+        - `images/` folder: Contains MRI scans
+        - `labels/` folder: Contains bounding box annotations
+        
+        **Converted to**: Binary Classification
+        - **No Tumor**: No bounding boxes in labels
+        - **Tumor**: One or more bounding boxes present
+        
+        **Splits**: test, valid
         """)
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()