Update app.py

app.py

@@ -6,12 +6,16 @@ from torchvision import models
 from transformers import BertTokenizer, BertModel
 import pandas as pd
 from datasets import load_dataset
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, random_split
 from sklearn.preprocessing import LabelEncoder
+from sklearn.metrics import confusion_matrix, classification_report, accuracy_score
+import seaborn as sns
+import matplotlib.pyplot as plt
+import numpy as np
+from tqdm import tqdm
 
 # Load dataset and filter out null/none values
 dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
-# Filter out entries where Model is None or empty
 dataset = dataset.filter(lambda example: example['Model'] is not None and example['Model'].strip() != '')
 
 # Preprocess text data
@@ -26,6 +30,9 @@ class CustomDataset(Dataset):
         ])
         self.label_encoder = LabelEncoder()
         self.labels = self.label_encoder.fit_transform(dataset['Model'])
+        
+        # Save unique model names for later use
+        self.unique_models = self.label_encoder.classes_
 
     def __len__(self):
         return len(self.dataset)
@@ -41,69 +48,183 @@ class CustomDataset(Dataset):
         label = self.labels[idx]
         return image, text, label
 
-# Define CNN for image processing
-class ImageModel(nn.Module):
-    def __init__(self):
-        super(ImageModel, self).__init__()
-        self.model = models.resnet18(pretrained=True)
-        self.model.fc = nn.Linear(self.model.fc.in_features, 512)
+# Model classes remain the same as before
+# ... (ImageModel, TextModel, CombinedModel classes stay unchanged)
 
-    def forward(self, x):
-        return self.model(x)
-
-# Define MLP for text processing
-class TextModel(nn.Module):
-    def __init__(self):
-        super(TextModel, self).__init__()
-        self.bert = BertModel.from_pretrained('bert-base-uncased')
-        self.fc = nn.Linear(768, 512)
+class ModelTrainerEvaluator:
+    def __init__(self, model, dataset, batch_size=32, learning_rate=0.001):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model = model.to(self.device)
+        self.batch_size = batch_size
+        self.criterion = nn.CrossEntropyLoss()
+        self.optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+        
+        # Split dataset into train, validation, and test
+        total_size = len(dataset)
+        train_size = int(0.7 * total_size)
+        val_size = int(0.15 * total_size)
+        test_size = total_size - train_size - val_size
+        
+        train_dataset, val_dataset, test_dataset = random_split(
+            dataset, [train_size, val_size, test_size]
+        )
+        
+        self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+        self.val_loader = DataLoader(val_dataset, batch_size=batch_size)
+        self.test_loader = DataLoader(test_dataset, batch_size=batch_size)
+        
+        self.unique_models = dataset.unique_models
 
-    def forward(self, x):
-        output = self.bert(**x)
-        return self.fc(output.pooler_output)
+    def train_epoch(self):
+        self.model.train()
+        total_loss = 0
+        predictions = []
+        actual_labels = []
+        
+        for batch in tqdm(self.train_loader, desc="Training"):
+            images, texts, labels = batch
+            images = images.to(self.device)
+            labels = labels.to(self.device)
+            
+            # Forward pass
+            self.optimizer.zero_grad()
+            outputs = self.model(images, texts)
+            loss = self.criterion(outputs, labels)
+            
+            # Backward pass
+            loss.backward()
+            self.optimizer.step()
+            
+            total_loss += loss.item()
+            
+            # Store predictions
+            _, preds = torch.max(outputs, 1)
+            predictions.extend(preds.cpu().numpy())
+            actual_labels.extend(labels.cpu().numpy())
+        
+        return total_loss / len(self.train_loader), predictions, actual_labels
 
-# Combined model
-class CombinedModel(nn.Module):
-    def __init__(self):
-        super(CombinedModel, self).__init__()
-        self.image_model = ImageModel()
-        self.text_model = TextModel()
-        self.fc = nn.Linear(1024, len(dataset['Model']))
+    def evaluate(self, loader, mode="Validation"):
+        self.model.eval()
+        total_loss = 0
+        predictions = []
+        actual_labels = []
+        
+        with torch.no_grad():
+            for batch in tqdm(loader, desc=mode):
+                images, texts, labels = batch
+                images = images.to(self.device)
+                labels = labels.to(self.device)
+                
+                outputs = self.model(images, texts)
+                loss = self.criterion(outputs, labels)
+                
+                total_loss += loss.item()
+                
+                _, preds = torch.max(outputs, 1)
+                predictions.extend(preds.cpu().numpy())
+                actual_labels.extend(labels.cpu().numpy())
+        
+        return total_loss / len(loader), predictions, actual_labels
 
-    def forward(self, image, text):
-        image_features = self.image_model(image)
-        text_features = self.text_model(text)
-        combined = torch.cat((image_features, text_features), dim=1)
-        return self.fc(combined)
+    def plot_confusion_matrix(self, y_true, y_pred, title):
+        cm = confusion_matrix(y_true, y_pred)
+        plt.figure(figsize=(15, 15))
+        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
+        plt.title(title)
+        plt.ylabel('True Label')
+        plt.xlabel('Predicted Label')
+        plt.savefig(f'{title.lower().replace(" ", "_")}.png')
+        plt.close()
 
-# Instantiate model
-model = CombinedModel()
+    def generate_evaluation_report(self, y_true, y_pred, title):
+        report = classification_report(y_true, y_pred, 
+                                    target_names=self.unique_models,
+                                    output_dict=True)
+        df_report = pd.DataFrame(report).transpose()
+        df_report.to_csv(f'{title.lower().replace(" ", "_")}_report.csv')
+        
+        accuracy = accuracy_score(y_true, y_pred)
+        
+        print(f"\n{title} Results:")
+        print(f"Accuracy: {accuracy:.4f}")
+        print("\nClassification Report:")
+        print(classification_report(y_true, y_pred, target_names=self.unique_models))
+        
+        return accuracy, df_report
 
-# Define predict function
-def predict(image):
-    model.eval()
-    with torch.no_grad():
-        image = transforms.ToTensor()(image).unsqueeze(0)
-        image = transforms.Resize((224, 224))(image)
-        text_input = tokenizer(
-            "Sample prompt",
-            return_tensors='pt',
-            padding=True,
-            truncation=True
-        )
-        output = model(image, text_input)
-        _, indices = torch.topk(output, 5)
-        recommended_models = [dataset['Model'][i] for i in indices[0]]
-    return recommended_models
+    def train_and_evaluate(self, num_epochs=5):
+        best_val_loss = float('inf')
+        train_accuracies = []
+        val_accuracies = []
+        
+        for epoch in range(num_epochs):
+            print(f"\nEpoch {epoch+1}/{num_epochs}")
+            
+            # Training
+            train_loss, train_preds, train_labels = self.train_epoch()
+            train_accuracy, _ = self.generate_evaluation_report(
+                train_labels, train_preds, f"Training Epoch {epoch+1}"
+            )
+            self.plot_confusion_matrix(
+                train_labels, train_preds, f"Training Confusion Matrix Epoch {epoch+1}"
+            )
+            
+            # Validation
+            val_loss, val_preds, val_labels = self.evaluate(self.val_loader)
+            val_accuracy, _ = self.generate_evaluation_report(
+                val_labels, val_preds, f"Validation Epoch {epoch+1}"
+            )
+            self.plot_confusion_matrix(
+                val_labels, val_preds, f"Validation Confusion Matrix Epoch {epoch+1}"
+            )
+            
+            train_accuracies.append(train_accuracy)
+            val_accuracies.append(val_accuracy)
+            
+            print(f"\nTraining Loss: {train_loss:.4f}")
+            print(f"Validation Loss: {val_loss:.4f}")
+            
+            # Save best model
+            if val_loss < best_val_loss:
+                best_val_loss = val_loss
+                torch.save(self.model.state_dict(), 'best_model.pth')
+        
+        # Plot training history
+        plt.figure(figsize=(10, 6))
+        plt.plot(train_accuracies, label='Training Accuracy')
+        plt.plot(val_accuracies, label='Validation Accuracy')
+        plt.title('Model Accuracy over Epochs')
+        plt.xlabel('Epoch')
+        plt.ylabel('Accuracy')
+        plt.legend()
+        plt.savefig('training_history.png')
+        plt.close()
+        
+        # Final test evaluation
+        self.model.load_state_dict(torch.load('best_model.pth'))
+        test_loss, test_preds, test_labels = self.evaluate(self.test_loader, "Test")
+        self.generate_evaluation_report(test_labels, test_preds, "Final Test")
+        self.plot_confusion_matrix(test_labels, test_preds, "Final Test Confusion Matrix")
 
-# Set up Gradio interface
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(label="Recommended Models"),
-    title="AI Image Model Recommender",
-    description="Upload an AI-generated image to receive model recommendations."
-)
+# Usage example
+def main():
+    # Create dataset
+    custom_dataset = CustomDataset(dataset)
+    
+    # Create model
+    model = CombinedModel()
+    
+    # Create trainer/evaluator
+    trainer = ModelTrainerEvaluator(
+        model=model,
+        dataset=custom_dataset,
+        batch_size=32,
+        learning_rate=0.001
+    )
+    
+    # Train and evaluate
+    trainer.train_and_evaluate(num_epochs=5)
 
-# Launch the app
-interface.launch()
+if __name__ == "__main__":
+    main()