Update app.py

app.py

@@ -8,11 +8,16 @@ import pandas as pd
 from datasets import load_dataset
 from torch.utils.data import DataLoader, Dataset
 from sklearn.preprocessing import LabelEncoder
+from sklearn.metrics import confusion_matrix, classification_report
+import seaborn as sns
+import matplotlib.pyplot as plt
+import numpy as np
 
-# Load dataset and filter out null/none values
+# Load dataset
 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() != '')
+
+# Filter out entries with None or null Model values
+filtered_dataset = dataset.filter(lambda example: example['Model'] is not None)
 
 # Preprocess text data
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -64,11 +69,11 @@ class TextModel(nn.Module):
 
 # Combined model
 class CombinedModel(nn.Module):
-    def __init__(self):
+    def __init__(self, num_classes):
         super(CombinedModel, self).__init__()
         self.image_model = ImageModel()
         self.text_model = TextModel()
-        self.fc = nn.Linear(1024, len(dataset['Model']))
+        self.fc = nn.Linear(1024, num_classes)
 
     def forward(self, image, text):
         image_features = self.image_model(image)
@@ -76,8 +81,42 @@ class CombinedModel(nn.Module):
         combined = torch.cat((image_features, text_features), dim=1)
         return self.fc(combined)
 
+def evaluate_model(model, test_loader, device):
+    model.eval()
+    all_preds = []
+    all_labels = []
+    
+    with torch.no_grad():
+        for images, texts, labels in test_loader:
+            images = images.to(device)
+            texts = {k: v.to(device) for k, v in texts.items()}
+            labels = labels.to(device)
+            
+            outputs = model(images, texts)
+            _, predicted = torch.max(outputs.data, 1)
+            
+            all_preds.extend(predicted.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+    
+    # Generate confusion matrix
+    cm = confusion_matrix(all_labels, all_preds)
+    
+    # Plot confusion matrix
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(cm, annot=True, fmt='d')
+    plt.title('Confusion Matrix')
+    plt.ylabel('True Label')
+    plt.xlabel('Predicted Label')
+    plt.savefig('confusion_matrix.png')
+    plt.close()
+    
+    # Print classification report
+    print(classification_report(all_labels, all_preds))
+
 # Instantiate model
-model = CombinedModel()
+dataset = CustomDataset(filtered_dataset)
+num_classes = len(np.unique(dataset.labels))
+model = CombinedModel(num_classes)
 
 # Define predict function
 def predict(image):
@@ -93,7 +132,7 @@ def predict(image):
         )
         output = model(image, text_input)
         _, indices = torch.topk(output, 5)
-        recommended_models = [dataset['Model'][i] for i in indices[0]]
+        recommended_models = [dataset.label_encoder.inverse_transform([i])[0] for i in indices[0]]
     return recommended_models
 
 # Set up Gradio interface
@@ -105,5 +144,6 @@ interface = gr.Interface(
     description="Upload an AI-generated image to receive model recommendations."
 )
 
-# Launch the app
-interface.launch()
+if __name__ == "__main__":
+    # Launch the app
+    interface.launch()