idek

main.py

@@ -3,9 +3,8 @@ import cv2
 import torch
 import numpy as np
 
-model = torch.load("models/bayes_cat_dog_classifier.pth")
+model = torch.load("bayes_cat_dog_classifier.pth")
 model.eval()
-
 model.to("cuda")
 
 IMG_SIZE = 128
@@ -24,11 +23,22 @@ def predict_image(image_path):
     return "Dog" if predicted == 1 else "Cat"
 
 
+# Cats
 preds = []
-for filename in os.listdir("dataset/test_set/XD/"):
-    img_path = os.path.join("dataset/test_set/XD/", filename)
+for filename in os.listdir("dataset/test_set/cats/"):
+    img_path = os.path.join("dataset/test_set/cats/", 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,8 +7,6 @@ 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")
@@ -28,6 +26,7 @@ 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)
@@ -51,9 +50,10 @@ if __name__ == "__main__":
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
 
-    num_epochs = 25
+    num_epochs = 20
     start_time = time.time()
 
+    model.train()
     for epoch in range(num_epochs):
         total_loss = 0
         for images, labels in dataloader:
@@ -67,4 +67,5 @@ 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"models/bayes_cat_dog_classifier.pth")
+
+    torch.save(model, f"bayes_cat_dog_classifier.pth")