README.md

@@ -5,9 +5,9 @@ This project is meant as a way to gradually bring improvements on the bayesian n
 ## Table of Contents
 
 * [Objectives 🎯](#objectives)
-* [Requirements 📋](#requirements)
 * [Running the project 🚀](#running-the-project)
 * [Development 🔨](#development)
+    * [Requirements 📋](#requirements)
 * [Citations 📝](#citations)
 
 ## Objectives 🎯
@@ -16,12 +16,6 @@ This project is meant as a way to gradually bring improvements on the bayesian n
 - [ ] Generate some graphs to visualize the data
 - [ ] Make a CLI
 
-## Requirements 📋
-
-To run the projet you need the following requirements:
-- Python 3.12
-- venv
-
 ## Running the project 🚀
 
 ```sh
@@ -35,4 +29,6 @@ python main.py
 
 ## Development 🔨
 
+### Requirements 📋
+
 ## Citations 📝

main.py

@@ -1,54 +1,46 @@
-import inquirer
-import typer
-import pyfiglet
-from yaspin import yaspin
-from train import train_model
-from predict import make_predictions
 import os
+import cv2
+import torch
+import numpy as np
 
-choice = ""
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-def main():
-    choice = inquirer.list_input("What would you like to do?", choices=["Run tests", "Train a model", "Visualize training data"])
-    
-    if choice == "Run tests":
-        predictions()
-    elif choice == "Train a model":
-        training()
-    else:
-        visualize()
+model = torch.load("bayes_cat_dog_classifier.pth")
+model.eval()
+model.to(DEVICE)
 
-def predictions():
-    default_cats_path = "dataset/test_set/cats/"
-    default_dogs_path = "dataset/test_set/dogs/"
-    models_base_path = "models/"
+IMG_SIZE = 128
 
-    model_name = inquirer.list_input("Select the model to use", choices=os.listdir(models_base_path))
-    model_path = os.path.join(models_base_path, model_name)
+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
 
-    dataset = inquirer.list_input("Select the testing data (default dataset)", choices=['Cats', 'Dogs'])
-    
-    if dataset == "Cats":
-        with yaspin(text="Making predictions...", color="cyan") as sp:
-            make_predictions(model_path, default_cats_path, "Cat", sp)
-            sp.ok("DONE")
-    else: 
-        with yaspin(text="Making predictions...", color="cyan") as sp:
-            make_predictions(model_path, default_dogs_path, "Dog", sp)
-            sp.ok("DONE")
+    model.eval()  # Set model to evaluation mode
+    with torch.no_grad():
+        output = model(img_tensor)
+        predicted = torch.argmax(output, dim=1).item()
 
-def training():
-    text = inquirer.text(message="Enter the name of the new model")
-    with yaspin(text="Training new model...", color="cyan") as sp:
-        train_model(text, sp)
-        sp.ok("DONE")
+    return "Dog" if predicted == 1 else "Cat"
 
-    
-def visualize():
-    print("Not available yet...\n")
-    main()
 
-if __name__ == "__main__":
-    print(pyfiglet.figlet_format("Cats and Dogs"))
-    print(pyfiglet.figlet_format("classification"))
-    typer.run(main)
+# 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)

predict.py

@@ -1,39 +0,0 @@
-import os
-import cv2
-import torch
-import numpy as np
-
-IMG_SIZE = 128
-
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-elif torch.mps.is_available():
-    DEVICE = torch.device("mps")
-else:
-    DEIVCE = torch.device("cpu")
-
-def predict_image(image_path, model):
-    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()
-    with torch.no_grad():
-        output = model(img_tensor)
-        predicted = torch.argmax(output, dim=1).item()
-
-    return "Dog" if predicted == 1 else "Cat"
-
-def make_predictions(model_path, dataset_path, type, spinner):
-    model = torch.load(model_path, map_location=DEVICE)
-    model.eval()
-    model.to(DEVICE)
-
-    preds = []
-    for filename in os.listdir(dataset_path):
-        img_path = os.path.join(dataset_path, filename)
-        prediction = predict_image(img_path, model)
-        preds.append(prediction)
-
-    spinner.write(f'Precision : {preds.count(type) / 1000 * 100}%')

requirements.txt

@@ -1,9 +1,5 @@
-inquirer==3.4.0
 matplotlib==3.10.0
 numpy==2.2.2
 opencv-python==4.11.0.86
-pyfiglet==1.0.2
 torch==2.5.1
-torchvision==0.20.1
-typer==0.15.1
-yaspin==3.1.0
+torchvision==0.20.1

train.py

@@ -6,17 +6,16 @@ from torch.utils.data import DataLoader, TensorDataset
 from torch import nn, optim
 from BN import CatDogClassifier
 import time
-from yaspin import yaspin
 
 IMG_SIZE = 128
 
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-elif torch.mps.is_available():
-    DEVICE = torch.device("mps")
-else:
-    DEIVCE = torch.device("cpu")
+#DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+DEVICE = "mps"
 
+
+"""
+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):
@@ -28,9 +27,7 @@ def load_images_from_folder(folder, label):
             data.append((img, label))
     return data
 
-def train_model(model_name, spinner):
-
-    spinner.write(f'Using the following device : {DEVICE}')
+if __name__ == "__main__":
 
     # Loading the dataset
     cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
@@ -69,8 +66,8 @@ def train_model(model_name, spinner):
             optimizer.step()
             total_loss += loss.item()
 
-        spinner.write(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}")
 
-    spinner.write(f"Time taken: {(time.time() - start_time):.2f} seconds")
+    print(f"Time taken: {(time.time() - start_time):.2f} seconds")
 
-    torch.save(model, f"models/{model_name}.pth")
+    torch.save(model, f"bayes_cat_dog_classifier.pth")