README.md

@@ -12,10 +12,9 @@ This project is meant as a way to gradually bring improvements on the bayesian n
 
 ## Objectives 🎯
 
-- [ ] Refactor the code
+- [ ] Pretrain some models
 - [ ] Generate some graphs to visualize the data
-- [X] Make a CLI
+- [ ] Make a CLI
-- [X] Pretrain some models
 
 ## Requirements 📋
 
@@ -26,11 +25,12 @@ To run the projet you need the following requirements:
 ## Running the project 🚀
 
 ```sh
-$ python -m venv .venv
+python -m venv .venv
-$ source .venv/bin/activate
+source .venv/bin/activate
-$ pip install requirements.txt
+pip install requirements.txt
 
-$ python main.py
+python train.py # If you want to train your own model
+python main.py
 ```
 
 ## Development 🔨

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():
+model = torch.load("bayes_cat_dog_classifier.pth")
-    choice = inquirer.list_input("What would you like to do?", choices=["Run tests", "Train a model", "Visualize training data"])
+model.eval()
-    
+model.to(DEVICE)
-    if choice == "Run tests":
-        predictions()
-    elif choice == "Train a model":
-        training()
-    else:
-        visualize()
 
-def predictions():
+IMG_SIZE = 128
-    default_cats_path = "dataset/test_set/cats/"
-    default_dogs_path = "dataset/test_set/dogs/"
-    models_base_path = "models/"
 
-    model_name = inquirer.list_input("Select the model to use", choices=os.listdir(models_base_path))
+def predict_image(image_path):
-    model_path = os.path.join(models_base_path, model_name)
+    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'])
+    model.eval()  # Set model to evaluation mode
-    
+    with torch.no_grad():
-    if dataset == "Cats":
+        output = model(img_tensor)
-        with yaspin(text="Making predictions...", color="cyan") as sp:
+        predicted = torch.argmax(output, dim=1).item()
-            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")
 
-def training():
+    return "Dog" if predicted == 1 else "Cat"
-    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")
 
-    
-def visualize():
-    print("Not available yet...\n")
-    main()
 
-if __name__ == "__main__":
+# Cats
-    print(pyfiglet.figlet_format("Cats and Dogs"))
+preds = []
-    print(pyfiglet.figlet_format("classification"))
+for filename in os.listdir("dataset/test_set/cats/"):
-    typer.run(main)
+    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:
-    DEVICE = 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,36 @@
-inquirer==3.4.0
+contourpy==1.3.1
+cycler==0.12.1
+filelock==3.17.0
+fonttools==4.55.6
+fsspec==2024.12.0
+Jinja2==3.1.5
+kiwisolver==1.4.8
+MarkupSafe==3.0.2
 matplotlib==3.10.0
+mpmath==1.3.0
+networkx==3.4.2
 numpy==2.2.2
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
 opencv-python==4.11.0.86
-pyfiglet==1.0.2
+packaging==24.2
+pillow==11.1.0
+pyparsing==3.2.1
+python-dateutil==2.9.0.post0
+setuptools==75.8.0
+six==1.17.0
+sympy==1.13.1
 torch==2.5.1
 torchvision==0.20.1
-typer==0.15.1
+triton==3.1.0
-yaspin==3.1.0
+typing_extensions==4.12.2

train.py

@@ -6,17 +6,14 @@ 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:0" if torch.cuda.is_available() else "cpu")
-    DEVICE = torch.device("cuda")
-elif torch.mps.is_available():
-    DEVICE = torch.device("mps")
-else:
-    DEVICE = torch.device("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):
@@ -28,9 +25,7 @@ def load_images_from_folder(folder, label):
             data.append((img, label))
     return data
 
-def train_model(model_name, spinner):
+if __name__ == "__main__":
-
-    spinner.write(f'Using the following device : {DEVICE}')
 
     # Loading the dataset
     cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
@@ -50,7 +45,7 @@ def train_model(model_name, spinner):
     dataset = TensorDataset(X_tensor, Y_tensor)
     dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
 
-    model = CatDogClassifier(img_size=IMG_SIZE)
+    model = CatDogClassifier()
     model = model.to(DEVICE)
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
@@ -69,8 +64,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")