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 torch
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
model = torch.load("models/bayes_cat_dog_classifier.pth")
|
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||||
model.eval()
|
|
||||||
|
|
||||||
model.to("cuda")
|
model = torch.load("bayes_cat_dog_classifier.pth")
|
||||||
|
model.eval()
|
||||||
|
model.to(DEVICE)
|
||||||
|
|
||||||
IMG_SIZE = 128
|
IMG_SIZE = 128
|
||||||
|
|
||||||
@ -14,7 +15,7 @@ def predict_image(image_path):
|
|||||||
img = cv2.imread(image_path)
|
img = cv2.imread(image_path)
|
||||||
img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) / 255.0
|
img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) / 255.0
|
||||||
img = np.transpose(img, (2, 0, 1)) # Convert to (C, H, W)
|
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
|
model.eval() # Set model to evaluation mode
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
@ -24,11 +25,22 @@ def predict_image(image_path):
|
|||||||
return "Dog" if predicted == 1 else "Cat"
|
return "Dog" if predicted == 1 else "Cat"
|
||||||
|
|
||||||
|
|
||||||
|
# Cats
|
||||||
preds = []
|
preds = []
|
||||||
for filename in os.listdir("dataset/test_set/XD/"):
|
for filename in os.listdir("dataset/test_set/cats/"):
|
||||||
img_path = os.path.join("dataset/test_set/XD/", filename)
|
img_path = os.path.join("dataset/test_set/cats/", filename)
|
||||||
prediction = predict_image(img_path)
|
prediction = predict_image(img_path)
|
||||||
preds.append(prediction)
|
preds.append(prediction)
|
||||||
|
|
||||||
print(preds.count("Cat"))
|
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"))
|
||||||
|
print(preds.count("Dog") / 1000)
|
9
train.py
9
train.py
@ -7,8 +7,6 @@ from torch import nn, optim
|
|||||||
from BN import CatDogClassifier
|
from BN import CatDogClassifier
|
||||||
import time
|
import time
|
||||||
|
|
||||||
from main import total_loss, outputs
|
|
||||||
|
|
||||||
IMG_SIZE = 128
|
IMG_SIZE = 128
|
||||||
|
|
||||||
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
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
|
return data
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|
||||||
# Loading the dataset
|
# Loading the dataset
|
||||||
cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
|
cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
|
||||||
dog_data = load_images_from_folder("dataset/training_set/dogs", label=1)
|
dog_data = load_images_from_folder("dataset/training_set/dogs", label=1)
|
||||||
@ -51,9 +50,10 @@ if __name__ == "__main__":
|
|||||||
criterion = nn.CrossEntropyLoss()
|
criterion = nn.CrossEntropyLoss()
|
||||||
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
||||||
|
|
||||||
num_epochs = 25
|
num_epochs = 20
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
|
model.train()
|
||||||
for epoch in range(num_epochs):
|
for epoch in range(num_epochs):
|
||||||
total_loss = 0
|
total_loss = 0
|
||||||
for images, labels in dataloader:
|
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"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}")
|
||||||
|
|
||||||
print(f"Time taken: {(time.time() - start_time):.2f} seconds")
|
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