From 5ffda30beb012fcb44ce0c7f6f5991de8957aaa8 Mon Sep 17 00:00:00 2001 From: Niels Date: Mon, 26 Aug 2024 13:27:48 +0200 Subject: [PATCH] Add mixin --- README.md | 4 ++-- clip_count/models/clip_count_model.py | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index cd2754a..9e5f628 100644 --- a/README.md +++ b/README.md @@ -51,11 +51,11 @@ The code can be easily adapted to other models, if you will to utilize CLIP-COUN ###### Evaluation -Our evaluation is based both on CLIP-COUNT and YOLO. +Our evaluation is based both on CLIP-COUNT and YOLOS. For CLIP-COUNT setup, refer to previous section. -For YOLO setup, please refer to [YOLOv9 docs](https://docs.ultralytics.com/models/yolov9/). +For YOLOS setup, please refer to the [Hugging Face docs](https://huggingface.co/docs/transformers/en/model_doc/yolos). ## Run and Evaluate:

diff --git a/clip_count/models/clip_count_model.py b/clip_count/models/clip_count_model.py index fcdf83b..547fc11 100644 --- a/clip_count/models/clip_count_model.py +++ b/clip_count/models/clip_count_model.py @@ -13,7 +13,12 @@ import einops import functools import operator -class CLIPCount(nn.Module): + +from huggingface_hub import PyTorchModelHubMixin + + +class CLIPCount(nn.Module, PyTorchModelHubMixin, + repo_url="https://github.com/songrise/CLIP-Count/tree/main", pipeline_tag="zero-shot-image-classification", license="mit"): def __init__(self, fim_depth:int=4, fim_num_heads:int=8, mlp_ratio:float=4.,