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