CoolFace
Modelpublic

Zahaab/object-classification

sourceHugging Facemitupdated 2y agoView on Hugging Face
0likes
Model Card

Image Classifier

This repository contains a pre-trained PyTorch model, designed for classifying images into 10 categories: airplane, bird, car, cat, deer, dog, horse, monkey, ship, and truck. The model uses a Convolutional Neural Network (CNN) architecture and can classify images based on the categories below.

Model Overview

The model is a simple CNN classifier with two convolutional blocks followed by a fully connected layer. It was trained on an image dataset and can classify images into the following categories:

  • —0: Airplane
  • —1: Bird
  • —2: Car
  • —3: Cat
  • —4: Deer
  • —5: Dog
  • —6: Horse
  • —7: Monkey
  • —8: Ship
  • —9: Truck

Model Architecture

The model consists of the following layers:

  1. 1.Conv Block 1: Two convolutional layers with ReLU activations followed by max pooling.
  2. 2.Conv Block 2: Two more convolutional layers with ReLU activations and max pooling.
  3. 3.Fully Connected Classifier: A linear layer that maps the features to 10 output categories.

Here’s the architecture of the model:

python
class CNNV0(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=hidden_units*576, out_features=output_shape)
        )

    def forward(self, x):
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
        x = self.classifier(x)
        return x

Requirements

  • —Python 3.7 or higher
  • —PyTorch 1.8 or higher
  • —torchvision (for loading and preprocessing images)

Usage

  1. 1.Clone this repository and install dependencies:
bash
   git clone <repository-url>
   cd <repository-folder>
   pip install torch torchvision
  1. 1.Load and use the model in your Python script:
python
   import torch
   from torchvision import transforms
   from PIL import Image

   # Load the model
   model = torch.load('model_0.pth')
   model.eval()  # Set to evaluation mode

   # Load and preprocess the image
   transform = transforms.Compose([
       transforms.Resize((224, 224)),
       transforms.ToTensor(),
   ])
   img = Image.open('path_to_image.jpg')
   img = transform(img).view(1, 3, 224, 224)  # Reshape to (1, 3, 224, 224) for batch processing

   # Predict
   with torch.no_grad():
       output = model(img)
       _, predicted = torch.max(output, 1)
       print("Predicted Aircraft Type:", predicted.item())