Compare commits
3 Commits
082be54fb5
...
49fa3a1cab
Author | SHA1 | Date | |
---|---|---|---|
49fa3a1cab | |||
cb6c695712 | |||
73c9b6bead |
BIN
bayes_cat_dog_classifier.pth
Normal file
BIN
bayes_cat_dog_classifier.pth
Normal file
Binary file not shown.
Binary file not shown.
Before Width: | Height: | Size: 3.1 MiB |
Binary file not shown.
Before Width: | Height: | Size: 3.9 MiB |
24
main.py
24
main.py
@ -3,10 +3,11 @@ import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
model = torch.load("models/bayes_cat_dog_classifier.pth")
|
||||
model.eval()
|
||||
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model.to("cuda")
|
||||
model = torch.load("bayes_cat_dog_classifier.pth")
|
||||
model.eval()
|
||||
model.to(DEVICE)
|
||||
|
||||
IMG_SIZE = 128
|
||||
|
||||
@ -14,7 +15,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("cuda") # Add batch dimension
|
||||
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():
|
||||
@ -24,11 +25,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)
|
9
train.py
9
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")
|
Loading…
x
Reference in New Issue
Block a user