Edit model card

David YOLOS Model

This repository contains a custom YOLOS model fine-tuned on the Balloon Dataset for object detection tasks. The model was trained using the PyTorch Lightning framework and is available for inference and further fine-tuning.

Model Details

  • Model Architecture: YOLOS (You Only Look One-level Object Structure)
  • Base Model: hustvl/yolos-small
  • Training Framework: PyTorch Lightning
  • Dataset: Balloon Dataset
  • Number of Classes: 1 (Balloon)

Installation and Usage

Installation

You can install the necessary libraries using:

pip install transformers torch torchvision

Usage

You can load and use the model with the following code:

from transformers import AutoModelForObjectDetection, AutoFeatureExtractor
from PIL import Image
import torch

# Load model and feature extractor
model_name = "your-username/my-custom-yolos-model"
model = AutoModelForObjectDetection.from_pretrained(model_name)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)

# Load an image
image = Image.open("path/to/your/image.jpg")

# Preprocess the image
inputs = feature_extractor(images=image, return_tensors="pt")
pixel_values = inputs['pixel_values']

# Perform inference
model.eval()
with torch.no_grad():
    outputs = model(pixel_values=pixel_values)

# Visualize the results
# (Insert visualization code here)

Model Performance

  • Training Loss: 0.0614
  • Validation Loss: 0.1784
  • Training Dataset: Balloon Dataset (61 images)
  • Validation Dataset: Balloon Dataset (13 images)
  • Number of Epochs: 18

Citation

If you use this model in your research, please cite:

Copy code
@misc{my-custom-yolos-model,
  author = {Your Name},
  title = {YOLOS Fine-tuned on Balloon Dataset},
  year = {2024},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/your-username/my-custom-yolos-model}},
}

License

This model is licensed under the MIT License. Feel free to use, modify, and distribute it as you see fit.

Copy code

You can copy and paste this Markdown into your README file on Hugging Face.

Downloads last month
32
Safetensors
Model size
30.7M params
Tensor type
F32
·
Inference Examples
Unable to determine this model's library. Check the docs .