aircraft-classifier / README.md
Zahaab's picture
Update README.md
044074c verified
|
raw
history blame
1.74 kB
metadata
license: mit
datasets:
  - Voxel51/FGVC-Aircraft
base_model:
  - timm/tf_efficientnet_b2.in1k
pipeline_tag: image-classification
tags:
  - aircraft
  - airplane

Aircraft Classifier

This repository contains a pre-trained PyTorch model for classifying aircraft types based on images. The model file aircraft_classifier.pth can be downloaded and used to classify images of various aircraft models.

Model Overview

The aircraft_classifier.pth file is a PyTorch model trained on a dataset of aircraft images. It achieves a test accuracy of 75.26% on the FGVC Aircraft test dataset, making it a reliable choice for identifying aircraft types. The model is designed to be lightweight and efficient for real-time applications.

Requirements

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

Usage

  1. Clone this repository and install dependencies.
    git clone <repository-url>
    cd <repository-folder>
    pip install torch torchvision
    
  2. Load and use the model in your Python script:
    import torch
    from torchvision import transforms
    from PIL import Image
    
    # Load the model
    model = torch.load('aircraft_classifier.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())