SceneViT-Nano-P8-64

A Nano Vision Transformer (ViT) model built from scratch for scene classification tasks, trained on the Intel Image Classification dataset.

Build and Train pipeline here:

Model Specifications

• Patch Size: 8x8
• Image Size: 224x224
• Hidden Size: 64
• Attention Heads: 4
• Hidden Layers: 8
• Validation Accuracy: ~84.30%

Training Summary & Progression

The custom Vision Transformer model was trained for 30 epochs on the Intel Image Classification dataset using a Cosine Annealing learning rate scheduler and early stopping.

Training and Validation Metrics Table

Training Curves Visualization

The graphs below illustrate the progression of Loss and Accuracy across all 30 training epochs:

The model achieves stable convergence with a peak validation accuracy of 84.30% and minimal overfitting.

Final Evaluation Results from TestSet

• Test Loss: 0.4381
• Test Accuracy: 83.90%
• Macro F1-Score: 83.95%

Usage

from PIL import Image
from torchvision import transforms
import torch.nn.functional as F
import matplotlib.pyplot as plt
import glob
import random
import torch
from transformers import AutoModelForImageClassification

# Step 1: Download and load the model from Hugging Face Hub first
print("Loading model from Hugging Face...")
model = AutoModelForImageClassification.from_pretrained(
    "TheSon2202/SceneViT-Nano-P8-64", trust_remote_code=True
)
model.eval()
print("Model loaded successfully!")


# Step 2: Define a clean prediction and visualization function using direct categories list
def predict_scene(
    model,
    image_search_path="/kaggle/input/datasets/thangcpd/intelimageclassification/data/seg_train/seg_train/**/*.jpg",
):
  # Get image size from the model's configuration
  image_size = getattr(model.config, "image_size", 224)

  # Define image preprocessing pipeline
  test_transform = transforms.Compose([
      transforms.Resize((image_size, image_size)),
      transforms.ToTensor(),
      transforms.Normalize(
          mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
      ),
  ])

  # Search for image paths in the dataset
  image_paths = glob.glob(image_search_path, recursive=True)
  if len(image_paths) == 0:
    image_paths = glob.glob("/kaggle/input/**/*.jpg", recursive=True)

  if len(image_paths) == 0:
    print("No image files found in the specified path.")
    return

  # Randomly pick an image and preprocess it
  img_path = random.choice(image_paths)
  raw_image = Image.open(img_path).convert("RGB")
  input_tensor = test_transform(raw_image).unsqueeze(0)

  # Run inference
  with torch.no_grad():
    outputs = model(input_tensor)
    logits = outputs.logits
    probabilities = F.softmax(logits, dim=-1)

    pred_idx = torch.argmax(probabilities, dim=-1).item()
    confidence = probabilities[0][pred_idx].item()

  # Directly use the fixed categories list matching model training
  categories = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
  predicted_label = (
      categories[pred_idx] if pred_idx < len(categories) else str(pred_idx)
  )

  # Display the image with result shown on the plot title
  plt.figure(figsize=(4, 4))
  plt.imshow(raw_image)
  plt.title(f"Pred: {predicted_label} ({confidence * 100:.1f}%)", color="green")
  plt.axis("off")
  plt.show()


# Example usage: Call the function using the pre-loaded model
predict_scene(model)

How to use ONNX

!pip install onnxruntime
!pip install onnxscript

onnx_file_path = hf_hub_download(
    repo_id="TheSon2202/SceneViT-Nano-P8-64", filename="vit_model.onnx"
)

# read model through ort_session
ort_session = ort.InferenceSession("vit_model.onnx")