Make sure Hugging Face download stats work, better discoverability

[NielsRogge]

Hi @LilianHollard,

Thanks for this nice work! I wrote a quick PoC to showcase that you can easily have integration with the 🤗 hub so that you can automatically load the various LeYOLO models using from_pretrained (and push them using push_to_hub), track download numbers for your models (similar to models in the Transformers library), and have nice model cards on a per-model basis. It leverages the PyTorchModelHubMixin class which allows to inherits these methods.

Usage is as follows:

from ultralytics import YOLO

model = YOLO.from_pretrained("your-hf-username/leyolo-nano")

# optionally, one can push a trained model to the hub
model.push_to_hub("nielsr/my-awesome-leyolo-model")

This means people don't need to manually download a checkpoint first in their local environment, it just loads automatically from the hub. We could move all checkpoints to separate repos to an organization on the hub or your personal user account if you're interested.

cc also @kadirnar who did a great job already at making the models available on 🤗. Would be great to leverage the PyTorchModelHubMixin from now on for new models as it makes downloads etc work :)

Would you be interested in this integration?

Kind regards,

Niels

[LilianHollard]

Hi @NielsRogge! I am sorry for the late response. I have a lot of work to do besides LeYOLO since I am still 100% focused on my PhD and other research, but I am attentive and continue to answer as quickly as possible.

Your pull request is very interesting indeed! Thank you for providing explanations. I'm finally going to motivate myself to create an account on HuggingFace.

I will try integrating PyTorchModelHubMixin locally and push your pull request if I am satisfied. I get back to you soon!

Thank you,

Lilian

[LilianHollard]

It should be working now! I added all LeYOLO weight on my HuggingFace model hub. Thank you for proposing HuggingFace integration.

Kind regards,

Lilian

[NielsRogge]

Hi @LilianHollard thanks so much for integrating with HF!

Btw, for YOLOv10 we also did some follow-up PRs which could also be done here, for instance:

• the PyTorchModelHubMixin class also allows to add metadata etc so that people who push a trained LeYOLO model to the hub will automatically get tagged with "leyolo", "object-detection" as pipeline tag etc.
• push_to_hub could be promoted so that people who train a LeYOLO model make it available for the open-source community
• the CLI could also be updated to leverage HF, see this commit as an example

[kadirnar]

@NielsRogge

Can you write sample inference code?

My Inference Code:

import gradio as gr
from ultralytics import YOLO
import spaces
import supervision as sv

BOX_ANNOTATOR = sv.BoxAnnotator()
LABEL_ANNOTATOR = sv.LabelAnnotator()

category_dict = {
    0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus',
    6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant',
    11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat',
    16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear',
    22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag',
    27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard',
    32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove',
    36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle',
    40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl',
    46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli',
    51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake',
    56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table',
    61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard',
    67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink',
    72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors',
    77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'
}

def attempt_download_from_hub(repo_id, hf_token=None):
    # https://github.com/fcakyon/yolov5-pip/blob/main/yolov5/utils/downloads.py
    from huggingface_hub import hf_hub_download, list_repo_files
    from huggingface_hub.utils._errors import RepositoryNotFoundError
    from huggingface_hub.utils._validators import HFValidationError
    try:
        repo_files = list_repo_files(repo_id=repo_id, repo_type='model', token=hf_token)
        model_file = [f for f in repo_files if f.endswith('.pt')][0]
        file = hf_hub_download(
            repo_id=repo_id,
            filename=model_file,
            repo_type='model',
            token=hf_token,
        )
        return file
    except (RepositoryNotFoundError, HFValidationError):
        return None

@spaces.GPU(duration=10)
def LeYOLO_inference(image, model_id, image_size, conf_threshold, iou_threshold):
    MODEL_PATH = attempt_download_from_hub(model_id)
    model = model = YOLO(MODEL_PATH)
    model.to('cuda')
    results = model(source=image, imgsz=image_size, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]
    detections = sv.Detections.from_ultralytics(results)

    labels = [
        f"{category_dict[class_id]} {confidence:.2f}"
        for class_id, confidence in zip(detections.class_id, detections.confidence)
    ]

    annotated_image = BOX_ANNOTATOR.annotate(
        scene=image, detections=detections)
    annotated_image = LABEL_ANNOTATOR.annotate(
        scene=annotated_image, detections=detections, labels=labels)

    return annotated_image