|
--- |
|
license: apache-2.0 |
|
language: |
|
- en |
|
pipeline_tag: object-detection |
|
tags: |
|
- code |
|
--- |
|
# David YOLOS Model |
|
|
|
This repository contains a custom YOLOS model fine-tuned on the [Balloon Dataset](https://github.com/matterport/Mask_RCNN/tree/master/samples/balloon) 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: |
|
|
|
```bash |
|
pip install transformers torch torchvision |
|
``` |
|
|
|
# Usage |
|
You can load and use the model with the following code: |
|
|
|
```python |
|
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: |
|
|
|
```bibtex |
|
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. |
