Update vit_model_test.py

vit_model_test.py

@@ -3,24 +3,28 @@ import torch.nn as nn
 from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms
 from transformers import ViTForImageClassification
+from PIL import Image
 import os
 import pandas as pd
 from sklearn.model_selection import train_test_split
-from sklearn.metrics import accuracy_score, precision_score, confusion_matrix, f1_score, average_precision_score, recall_score
+from sklearn.metrics import accuracy_score, precision_score, confusion_matrix, f1_score, average_precision_score
 import matplotlib.pyplot as plt
 import seaborn as sns
+from sklearn.metrics import recall_score
+from vit_model_traning import labeling,CustomDataset
 
-# פונקציה להחזרת HTML של סרטון
-def display_video(video_url):
-    return f'''
-    <iframe width="640" height="480" src="{video_url}" frameborder="0" allowfullscreen></iframe>
-    '''
 
+    
 def shuffle_and_split_data(dataframe, test_size=0.2, random_state=59):
+    # Shuffle the DataFrame
     shuffled_df = dataframe.sample(frac=1, random_state=random_state).reset_index(drop=True)
+    
+    # Split the DataFrame into train and validation sets
     train_df, val_df = train_test_split(shuffled_df, test_size=test_size, random_state=random_state)
+    
     return train_df, val_df    
 
+
 if __name__ == "__main__":
     # Check for GPU availability
     device = torch.device('cuda')
@@ -28,8 +32,9 @@ if __name__ == "__main__":
     # Load the pre-trained ViT model and move it to GPU
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(device)
 
+
+
     model.classifier = nn.Linear(model.config.hidden_size, 2).to(device)
-    
     # Define the image preprocessing pipeline
     preprocess = transforms.Compose([
         transforms.Resize((224, 224)),
@@ -52,20 +57,9 @@ if __name__ == "__main__":
     true_labels = []
     predicted_labels = []
 
-    # קישור לסרטון ביוטיוב
-    video_url = 'https://www.youtube.com/embed/vGRq060nPYU'  # החלף ב-URL של הסרטון שלך
-    video_html = display_video(video_url)
-
-    # הראי את הסרטון לפני החיזוי
-    print(video_html)  # זה אמור להציג את ה-HTML בדשבורד שלך
-
     with torch.no_grad():
         for images, labels in test_loader:
             images, labels = images.to(device), labels.to(device)
-
-            # הראה את הסרטון בעת חיזוי
-            print(video_html)  # הצג את ה-HTML של הסרטון
-
             outputs = model(images)
             logits = outputs.logits  # Extract logits from the output
             _, predicted = torch.max(logits, 1)
@@ -80,6 +74,7 @@ if __name__ == "__main__":
     ap = average_precision_score(true_labels, predicted_labels)
     recall = recall_score(true_labels, predicted_labels)
 
+
     print(f"Test Accuracy: {accuracy:.2%}")
     print(f"Precision: {precision:.2%}")
     print(f"F1 Score: {f1:.2%}")