You can deploy this model and start earning money today!

Zahaab/
object-classification

jxie/stl10

mit

Views

No views yet

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:

Conv Block 1: Two convolutional layers with ReLU activations followed by max pooling.
Conv Block 2: Two more convolutional layers with ReLU activations and max pooling.
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

Clone this repository and install dependencies:

bash
1git clone <repository-url> 2cd <repository-folder> 3pip install torch torchvision
Load and use the model in your Python script:

python
1import torch 2from torchvision import transforms 3from PIL import Image 4 5# Load the model 6model = torch.load('model_0.pth') 7model.eval() # Set to evaluation mode 8 9# Load and preprocess the image 10transform = transforms.Compose([ 11 transforms.Resize((224, 224)), 12 transforms.ToTensor(), 13]) 14img = Image.open('path_to_image.jpg') 15img = transform(img).view(1, 3, 224, 224) # Reshape to (1, 3, 224, 224) for batch processing 16 17# Predict 18with torch.no_grad(): 19 output = model(img) 20 _, predicted = torch.max(output, 1) 21 print("Predicted Aircraft Type:", predicted.item())

Zahaab/object-classification