import os
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.eval()
model.to(DEVICE)

IMG_SIZE = 128

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

    model.eval()  # Set model to evaluation mode
    with torch.no_grad():
        output = model(img_tensor)
        predicted = torch.argmax(output, dim=1).item()

    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)
    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)