--- 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.