diff --git a/.gitignore b/.gitignore index b694934..0e5ac79 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ -.venv \ No newline at end of file +.venv +__pycache__ \ No newline at end of file diff --git a/BN.py b/BN.py index 65e507d..79305ea 100644 --- a/BN.py +++ b/BN.py @@ -1,5 +1,4 @@ import torch.nn as nn -import torch.optim as optim # Define CNN Model class CatDogClassifier(nn.Module): @@ -8,14 +7,14 @@ class CatDogClassifier(nn.Module): self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.pool = nn.MaxPool2d(2, 2) - self.fc1 = nn.Linear(64 * (img_size//4) * (img_size//4), 128) # Flatten size depends on image size - self.fc2 = nn.Linear(128, 2) # Output 2 classes (cat or dog) + self.fc1 = nn.Linear(64 * (img_size//4) * (img_size//4), 128) + self.fc2 = nn.Linear(128, 2) self.relu = nn.ReLU() def forward(self, x): x = self.pool(self.relu(self.conv1(x))) x = self.pool(self.relu(self.conv2(x))) - x = x.view(x.size(0), -1) # Flatten + x = x.view(x.size(0), -1) # Flattening x = self.relu(self.fc1(x)) x = self.fc2(x) return x diff --git a/README.md b/README.md new file mode 100644 index 0000000..356eaae --- /dev/null +++ b/README.md @@ -0,0 +1,9 @@ +# Cats and Dogs Classification + +--- + +Quick and dirty *Bayesian Network* that classifies pictures of cats and dogs. Works okay-ish.. + + +## Installations + diff --git a/caca.py b/caca.py deleted file mode 100644 index 0bbf357..0000000 --- a/caca.py +++ /dev/null @@ -1,7 +0,0 @@ -import torch -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -tensor = torch.randn(3, 3).to(device) # Moves the tensor to GPU - - -print(tensor) -print(device) \ No newline at end of file diff --git a/dataset/test_set/XD/IMG_2763.jpg b/dataset/test_set/XD/IMG_2763.jpg new file mode 100644 index 0000000..5328f87 Binary files /dev/null and b/dataset/test_set/XD/IMG_2763.jpg differ diff --git a/dataset/test_set/XD/IMG_3077.jpg b/dataset/test_set/XD/IMG_3077.jpg new file mode 100644 index 0000000..9df0ea0 Binary files /dev/null and b/dataset/test_set/XD/IMG_3077.jpg differ diff --git a/main.py b/main.py index 80f50d2..ad60870 100644 --- a/main.py +++ b/main.py @@ -2,71 +2,33 @@ import os import cv2 import torch import numpy as np -from torch.utils.data import DataLoader, TensorDataset -from torch import nn, optim -from BN import CatDogClassifier -import time -# Define image size +model = torch.load("models/bayes_cat_dog_classifier.pth") +model.eval() + +model.to("cuda") + IMG_SIZE = 128 -# Function to load images from a directory -def load_images_from_folder(folder, label): - data = [] - for filename in os.listdir(folder): - img_path = os.path.join(folder, filename) - img = cv2.imread(img_path) # Load image - if img is not None: - img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) # Resize - img = img / 255.0 # Normalize - data.append((img, label)) - return data +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 -# Load dataset -cat_data = load_images_from_folder("dataset/training_set/cats", label=0) # Label 0 for cats -dog_data = load_images_from_folder("dataset/training_set/dogs", label=1) # Label 1 for dogs + model.eval() # Set model to evaluation mode + with torch.no_grad(): + output = model(img_tensor) + predicted = torch.argmax(output, dim=1).item() -# Combine and shuffle data -dataset = cat_data + dog_data -np.random.shuffle(dataset) - -# Convert to NumPy arrays -X = np.array([item[0] for item in dataset], dtype=np.float32) # Image data -Y = np.array([item[1] for item in dataset], dtype=np.int64) # Labels - -# Reshape to PyTorch format (N, C, H, W) -X = np.transpose(X, (0, 3, 1, 2)) # Convert to (batch, channels, height, width) - -# Convert to PyTorch tensors -X_tensor = torch.tensor(X).to("cuda") -Y_tensor = torch.tensor(Y).to("cuda") - -# Create dataset and data loader -dataset = TensorDataset(X_tensor, Y_tensor) -dataloader = DataLoader(dataset, batch_size=32, shuffle=True) - -# Initialize model, loss, and optimizer -model = CatDogClassifier() -model = model.to("cuda") -criterion = nn.CrossEntropyLoss() -optimizer = optim.Adam(model.parameters(), lr=0.001) - -# Training loop -num_epochs = 25 -start_time = time.time() -for epoch in range(num_epochs): - total_loss = 0 - for images, labels in dataloader: - optimizer.zero_grad() # Reset gradients - outputs = model(images) # Forward pass - loss = criterion(outputs, labels) # Compute loss - loss.backward() # Backpropagation - optimizer.step() # Update weights - total_loss += loss.item() - - print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}") - -print(f"Time taken: {(time.time()-start_time):.2f} seconds") + return "Dog" if predicted == 1 else "Cat" -torch.save(model, "super_model.pth") \ No newline at end of file +preds = [] +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("Dog")) diff --git a/predict.py b/predict.py deleted file mode 100644 index 9daf22c..0000000 --- a/predict.py +++ /dev/null @@ -1,38 +0,0 @@ -import os -import cv2 -import torch -import numpy as np -from torch.utils.data import DataLoader, TensorDataset -from torch import nn, optim -from BN import CatDogClassifier - -model = torch.load("super_model.pth") -model.eval() - -model.to("cuda") - -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("cuda") # 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" - - -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("Dog")) -print(preds.count("Cat") / 1000) diff --git a/super_model.pth b/super_model.pth deleted file mode 100644 index 1e0d059..0000000 Binary files a/super_model.pth and /dev/null differ diff --git a/train.py b/train.py new file mode 100644 index 0000000..d3acdb3 --- /dev/null +++ b/train.py @@ -0,0 +1,70 @@ +import os +import cv2 +import numpy as np +import torch +from torch.utils.data import DataLoader, TensorDataset +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") + +""" +This function loads all the images from the folder and labels them +""" +def load_images_from_folder(folder, label): + data = [] + for filename in os.listdir(folder): + img_path = os.path.join(folder, filename) + img = cv2.imread(img_path) + if img is not None: + img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) + img = img / 255.0 + data.append((img, 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) + + dataset = cat_data + dog_data + np.random.shuffle(dataset) + + X = np.array([item[0] for item in dataset], dtype=np.float32) + Y = np.array([item[1] for item in dataset], dtype=np.int64) + + X = np.transpose(X, (0, 3, 1, 2)) + + X_tensor = torch.tensor(X).to(DEVICE) + Y_tensor = torch.tensor(Y).to(DEVICE) + + dataset = TensorDataset(X_tensor, Y_tensor) + dataloader = DataLoader(dataset, batch_size=32, shuffle=True) + + model = CatDogClassifier() + model = model.to(DEVICE) + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=0.001) + + num_epochs = 25 + start_time = time.time() + + for epoch in range(num_epochs): + total_loss = 0 + for images, labels in dataloader: + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + total_loss += loss.item() + + 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") \ No newline at end of file