Initial commit

.gitignore

@@ -1 +1,2 @@
-.venv
+.venv
+__pycache__

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

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
+

caca.py

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

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

predict.py

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

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