Update vit_model_test.py

vit_model_test.py

@@ -10,7 +10,7 @@ class CustomModel:
         # Explicitly set the device to CPU
         self.device = torch.device('cpu')
 
-        # Load the pre-trained ViT model and move it to CPU
+        # Load the pre-trained ViT 
         self.model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(self.device)
         self.model.classifier = nn.Linear(self.model.config.hidden_size, 2).to(self.device)
 
@@ -18,7 +18,7 @@ class CustomModel:
         self.model.load_state_dict(torch.load('trained_model.pth', map_location=self.device, weights_only=True))
         self.model.eval()
 
-        # Define the image preprocessing pipeline
+        # Resize the image and make it a tensor (add dimension)
         self.preprocess = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor()
@@ -26,7 +26,7 @@ class CustomModel:
 
     def predict(self, image: Image.Image):
         # Preprocess the image
-        image = self.preprocess(image).unsqueeze(0).to(self.device)  # Add batch dimension
+        image = self.preprocess(image).unsqueeze(0).to(self.device)  
         
         # Perform inference
         with torch.no_grad():
@@ -35,6 +35,6 @@ class CustomModel:
             probabilities = F.softmax(logits, dim=1)
             confidences, predicted = torch.max(probabilities, 1)
             predicted_label = predicted.item()
-            confidence = confidences.item() * 100  # Convert to percentage
+            confidence = confidences.item() * 100  # Convert to percentage format
 
         return predicted_label, confidence