dev: uploaded model files

cat_dog_classifier.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6be6dd4dbb80eb2824563dd9237b63a582a57482130e0494642e4e06ece39728
+size 1685764

config.json

@@ -0,0 +1,8 @@
+{
+  "batch_size": 64,
+  "num_epochs": 5,
+  "learning_rate": 0.001,
+  "model_save_path": "cat_dog_classifier.bin",
+  "data_path": "quickdraw_data",
+  "num_samples": 5000
+}

model.py

@@ -0,0 +1,13 @@
+import torch.nn as nn
+import torch
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super(SimpleModel, self).__init__()
+        self.fc1 = nn.Linear(784, 128)
+        self.fc2 = nn.Linear(128, 2)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x

quickdraw_data/cat.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21a281839d3f2eef601d57d2338a4eafdf24649f8d0a0e42d3ec3e595911463e
+size 96590448

quickdraw_data/dog.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72f95508d440976a075e7098557647bbdeaea7a06c63889215c5b87fbf82ea2c
+size 119292736

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.0.0
+numpy==1.24.3
+requests==2.31.0
+Pillow==9.4.0
+matplotlib==3.8.0
+scikit-learn==1.3.0

train_cat_dog_classifier.py

@@ -0,0 +1,310 @@
+import os
+import json
+import numpy as np
+import torch
+from torch.utils.data import TensorDataset, DataLoader, random_split
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from PIL import Image
+import requests
+import matplotlib.pyplot as plt
+
+# Ensure that matplotlib does not try to open a window (useful if running on a server)
+import matplotlib
+matplotlib.use('Agg')
+
+# Check if CUDA is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f'Using device: {device}')
+
+def load_config(config_file='config.json'):
+    """
+    Loads configuration parameters from a JSON file.
+    Args:
+        config_file (str): Path to the JSON config file.
+    Returns:
+        config (dict): Dictionary containing configuration parameters.
+    """
+    with open(config_file, 'r') as f:
+        return json.load(f)
+
+def download_quickdraw_data():
+    """
+    Downloads 'cat.npy' and 'dog.npy' files from the Quick, Draw! dataset.
+    """
+    os.makedirs('quickdraw_data', exist_ok=True)
+    base_url = 'https://storage.googleapis.com/quickdraw_dataset/full/numpy_bitmap/'
+
+    categories = ['cat', 'dog']
+    for category in categories:
+        url = f"{base_url}{category}.npy"
+        save_path = os.path.join('quickdraw_data', f"{category}.npy")
+
+        if os.path.exists(save_path):
+            print(f"{category}.npy already exists, skipping download.")
+            continue
+
+        print(f"Downloading {category}.npy...")
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(save_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print(f"Downloaded {category}.npy")
+        else:
+            print(f"Failed to download {category}.npy. Status code: {response.status_code}")
+
+def load_and_preprocess_data(num_samples=5000):
+    """
+    Loads and preprocesses the data for 'cat' and 'dog' categories.
+    Args:
+        num_samples (int): Number of samples to load for each category.
+    Returns:
+        train_loader, test_loader: DataLoaders for training and testing.
+    """
+    # Load data
+    cat_data = np.load('quickdraw_data/cat.npy')
+    dog_data = np.load('quickdraw_data/dog.npy')
+
+    # Limit the number of samples
+    cat_data = cat_data[:num_samples]
+    dog_data = dog_data[:num_samples]
+
+    # Create labels: 0 for cat, 1 for dog
+    cat_labels = np.zeros(len(cat_data), dtype=np.int64)
+    dog_labels = np.ones(len(dog_data), dtype=np.int64)
+
+    # Combine data and labels
+    data = np.concatenate((cat_data, dog_data), axis=0)
+    labels = np.concatenate((cat_labels, dog_labels), axis=0)
+
+    # Normalize data
+    data = data.astype('float32') / 255.0
+
+    # Reshape data for PyTorch: (batch_size, channels, height, width)
+    data = data.reshape(-1, 1, 28, 28)
+
+    # Convert to PyTorch tensors
+    data_tensor = torch.tensor(data)
+    labels_tensor = torch.tensor(labels)
+
+    # Create a TensorDataset
+    dataset = TensorDataset(data_tensor, labels_tensor)
+
+    # Split dataset into training and testing sets (80% train, 20% test)
+    train_size = int(0.8 * len(dataset))
+    test_size = len(dataset) - train_size
+
+    train_dataset, test_dataset = random_split(dataset, [train_size, test_size])
+
+    # Create DataLoaders
+    config = load_config()
+    batch_size = config['batch_size']
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size)
+
+    return train_loader, test_loader
+
+class SimpleCNN(nn.Module):
+    """
+    Defines a simple Convolutional Neural Network for binary classification.
+    """
+    def __init__(self):
+        super(SimpleCNN, self).__init__()
+        # Convolutional layers
+        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        # Fully connected layers
+        self.fc1 = nn.Linear(64 * 7 * 7, 128)
+        self.fc2 = nn.Linear(128, 2)  # 2 output classes: cat and dog
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))  # Convolutional layer 1
+        x = self.pool(x)           # Max pooling
+        x = F.relu(self.conv2(x))  # Convolutional layer 2
+        x = self.pool(x)           # Max pooling
+        x = x.view(-1, 64 * 7 * 7) # Flatten
+        x = F.relu(self.fc1(x))    # Fully connected layer 1
+        x = self.fc2(x)            # Output layer
+        return x
+
+def train_model(model, train_loader, num_epochs=5, learning_rate=0.001):
+    """
+    Trains the model using the training DataLoader.
+    Args:
+        model: The neural network model to train.
+        train_loader: DataLoader for the training data.
+        num_epochs (int): Number of epochs to train.
+        learning_rate (float): Learning rate for the optimizer.
+    """
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+    model.train()
+    for epoch in range(num_epochs):
+        running_loss = 0.0
+
+        for images, labels in train_loader:
+            images = images.to(device)
+            labels = labels.to(device)
+
+            # Zero the parameter gradients
+            optimizer.zero_grad()
+
+            # Forward pass
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            # Backward pass and optimize
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item() * images.size(0)
+
+        epoch_loss = running_loss / len(train_loader.dataset)
+        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}')
+
+def evaluate_model(model, test_loader):
+    """
+    Evaluates the model on the test DataLoader.
+    Args:
+        model: The trained neural network model.
+        test_loader: DataLoader for the test data.
+    """
+    model.eval()
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+        for images, labels in test_loader:
+            images = images.to(device)
+            labels = labels.to(device)
+
+            outputs = model(images)
+            _, predicted = torch.max(outputs.data, 1)
+
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+    accuracy = 100 * correct / total
+    print(f'Test Accuracy: {accuracy:.2f}%')
+
+def save_model(model, filepath='cat_dog_classifier.pth'):
+    """
+    Saves the trained model to a file.
+    Args:
+        model: The trained neural network model.
+        filepath (str): The path where the model will be saved.
+    """
+    torch.save(model.state_dict(), filepath)
+    print(f'Model saved to {filepath}')
+
+def load_model(model, filepath='cat_dog_classifier.pth'):
+    """
+    Loads the model parameters from a file.
+    Args:
+        model: The neural network model to load parameters into.
+        filepath (str): The path to the saved model file.
+    """
+    model.load_state_dict(torch.load(filepath, map_location=device))
+    model.to(device)
+    print(f'Model loaded from {filepath}')
+
+def predict_image(model, image):
+    """
+    Predicts the class of a single image.
+    Args:
+        model: The trained neural network model.
+        image: A PIL Image or NumPy array.
+    Returns:
+        prediction (str): The predicted class label ('cat' or 'dog').
+    """
+    # Preprocess the image
+    if isinstance(image, Image.Image):
+        image = image.resize((28, 28)).convert('L')
+        image = np.array(image).astype('float32') / 255.0
+    elif isinstance(image, np.ndarray):
+        if image.shape != (28, 28):
+            image = Image.fromarray(image).resize((28, 28)).convert('L')
+            image = np.array(image).astype('float32') / 255.0
+    else:
+        raise ValueError("Image must be a PIL Image or NumPy array.")
+
+    image = image.reshape(1, 1, 28, 28)
+    image_tensor = torch.tensor(image).to(device)
+
+    # Get prediction
+    model.eval()
+    with torch.no_grad():
+        output = model(image_tensor)
+        _, predicted = torch.max(output.data, 1)
+    return 'cat' if predicted.item() == 0 else 'dog'
+
+def visualize_predictions(model, test_loader, num_images=8):
+    """
+    Visualizes sample predictions from the test set.
+    Args:
+        model: The trained neural network model.
+        test_loader: DataLoader for the test data.
+        num_images (int): Number of images to display.
+    """
+    model.eval()
+    dataiter = iter(test_loader)
+    images, labels = next(dataiter)  # Use the built-in next() function
+
+    images = images.to(device)
+    labels = labels.to(device)
+
+    # Get predictions
+    outputs = model(images)
+    _, predicted = torch.max(outputs, 1)
+
+    # Move images to CPU for plotting
+    images = images.cpu().numpy()
+    predicted = predicted.cpu().numpy()
+    labels = labels.cpu().numpy()
+
+    # Plot the images with predicted and true labels
+    fig = plt.figure(figsize=(10, 4))
+    for idx in range(num_images):
+        ax = fig.add_subplot(2, num_images // 2, idx+1)
+        img = images[idx][0]
+        ax.imshow(img, cmap='gray')
+        pred_label = 'cat' if predicted[idx] == 0 else 'dog'
+        true_label = 'cat' if labels[idx] == 0 else 'dog'
+        ax.set_title(f'Pred: {pred_label}\nTrue: {true_label}')
+        ax.axis('off')
+    plt.tight_layout()
+    plt.savefig('sample_predictions.png')
+    print('Sample predictions saved to sample_predictions.png')
+
+def main():
+    # Load configuration
+    config = load_config()
+
+    # Step 1: Download the data
+    download_quickdraw_data()
+
+    # Step 2: Load and preprocess the data
+    train_loader, test_loader = load_and_preprocess_data(num_samples=config['num_samples'])
+
+    # Step 3: Initialize the model
+    model = SimpleCNN().to(device)
+
+    # Step 4: Train the model
+    train_model(model, train_loader, num_epochs=config['num_epochs'], learning_rate=config['learning_rate'])
+
+    # Step 5: Evaluate the model
+    evaluate_model(model, test_loader)
+
+    # Step 6: Visualize sample predictions
+    visualize_predictions(model, test_loader, num_images=8)
+
+    # Step 7: Save the model
+    save_model(model, config['model_save_path'])
+
+if __name__ == '__main__':
+    main()