main.py

@@ -3,11 +3,10 @@ import cv2
 import torch
 import numpy as np
 
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-model = torch.load("bayes_cat_dog_classifier.pth")
+model = torch.load("models/bayes_cat_dog_classifier.pth")
 model.eval()
-model.to(DEVICE)
+
+model.to("cuda")
 
 IMG_SIZE = 128
 
@@ -15,7 +14,7 @@ def predict_image(image_path):
     img = cv2.imread(image_path)
     img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) / 255.0
     img = np.transpose(img, (2, 0, 1))  # Convert to (C, H, W)
-    img_tensor = torch.tensor(img, dtype=torch.float32).unsqueeze(0).to(DEVICE)  # Add batch dimension
+    img_tensor = torch.tensor(img, dtype=torch.float32).unsqueeze(0).to("cuda")  # Add batch dimension
 
     model.eval()  # Set model to evaluation mode
     with torch.no_grad():
@@ -25,22 +24,11 @@ def predict_image(image_path):
     return "Dog" if predicted == 1 else "Cat"
 
 
-# Cats
 preds = []
-for filename in os.listdir("dataset/test_set/cats/"):
-    img_path = os.path.join("dataset/test_set/cats/", filename)
+for filename in os.listdir("dataset/test_set/XD/"):
+    img_path = os.path.join("dataset/test_set/XD/", filename)
     prediction = predict_image(img_path)
     preds.append(prediction)
 
 print(preds.count("Cat"))
-print(preds.count("Cat") / 1000)
-
-# Dogs
-preds = []
-for filename in os.listdir("dataset/test_set/dogs/"):
-    img_path = os.path.join("dataset/test_set/dogs/", filename)
-    prediction = predict_image(img_path)
-    preds.append(prediction)
-
 print(preds.count("Dog"))
-print(preds.count("Dog") / 1000)

train.py

@@ -7,6 +7,8 @@ from torch import nn, optim
 from BN import CatDogClassifier
 import time
 
+from main import total_loss, outputs
+
 IMG_SIZE = 128
 
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -26,7 +28,6 @@ def load_images_from_folder(folder, label):
     return data
 
 if __name__ == "__main__":
-
     # Loading the dataset
     cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
     dog_data = load_images_from_folder("dataset/training_set/dogs", label=1)
@@ -50,10 +51,9 @@ if __name__ == "__main__":
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
 
-    num_epochs = 20
+    num_epochs = 25
     start_time = time.time()
 
-    model.train()
     for epoch in range(num_epochs):
         total_loss = 0
         for images, labels in dataloader:
@@ -67,5 +67,4 @@ if __name__ == "__main__":
         print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}")
 
     print(f"Time taken: {(time.time() - start_time):.2f} seconds")
-
-    torch.save(model, f"bayes_cat_dog_classifier.pth")
+    torch.save(model, f"models/bayes_cat_dog_classifier.pth")