pre trained 2 models

README.md

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

main.py

@@ -1,46 +1,54 @@
+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
 
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+choice = ""
 
-model = torch.load("bayes_cat_dog_classifier.pth")
-model.eval()
-model.to(DEVICE)
+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()
 
-IMG_SIZE = 128
+def predictions():
+    default_cats_path = "dataset/test_set/cats/"
+    default_dogs_path = "dataset/test_set/dogs/"
+    models_base_path = "models/"
 
-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
+    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)
 
-    model.eval()  # Set model to evaluation mode
-    with torch.no_grad():
-        output = model(img_tensor)
-        predicted = torch.argmax(output, dim=1).item()
+    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")
 
-    return "Dog" if predicted == 1 else "Cat"
+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")
 
+    
+def visualize():
+    print("Not available yet...\n")
+    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)
+if __name__ == "__main__":
+    print(pyfiglet.figlet_format("Cats and Dogs"))
+    print(pyfiglet.figlet_format("classification"))
+    typer.run(main)

predict.py

@@ -0,0 +1,39 @@
+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,36 +1,9 @@
-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
+inquirer==3.4.0
 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
-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
+pyfiglet==1.0.2
 torch==2.5.1
 torchvision==0.20.1
-triton==3.1.0
-typing_extensions==4.12.2
+typer==0.15.1
+yaspin==3.1.0

train.py

@@ -6,14 +6,17 @@ 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
 
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+if torch.cuda.is_available():
+    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):
@@ -25,7 +28,9 @@ def load_images_from_folder(folder, label):
             data.append((img, label))
     return data
 
-if __name__ == "__main__":
+def train_model(model_name, spinner):
+
+    spinner.write(f'Using the following device : {DEVICE}')
 
     # Loading the dataset
     cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
@@ -45,7 +50,7 @@ if __name__ == "__main__":
     dataset = TensorDataset(X_tensor, Y_tensor)
     dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
 
-    model = CatDogClassifier()
+    model = CatDogClassifier(img_size=IMG_SIZE)
     model = model.to(DEVICE)
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
@@ -64,8 +69,8 @@ if __name__ == "__main__":
             optimizer.step()
             total_loss += loss.item()
 
-        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}")
+        spinner.write(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}")
 
-    print(f"Time taken: {(time.time() - start_time):.2f} seconds")
+    spinner.write(f"Time taken: {(time.time() - start_time):.2f} seconds")
 
-    torch.save(model, f"bayes_cat_dog_classifier.pth")
+    torch.save(model, f"models/{model_name}.pth")