Adding licensing information to cspnet.py

Author: alexanderdann

timm/models/_hub.py

@@ -30,8 +30,8 @@
 from timm.models._pretrained import filter_pretrained_cfg
 
 try:
-    from huggingface_hub import HfApi, hf_hub_download
-    from huggingface_hub.utils import EntryNotFoundError
+    from huggingface_hub import HfApi, hf_hub_download, model_info
+    from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
     hf_hub_download = partial(hf_hub_download, library_name="timm", library_version=__version__)
     _has_hf_hub = True
 except ImportError:
@@ -533,3 +533,44 @@ def _get_safe_alternatives(filename: str) -> Iterable[str]:
         yield HF_OPEN_CLIP_SAFE_WEIGHTS_NAME
     if filename not in (HF_WEIGHTS_NAME, HF_OPEN_CLIP_WEIGHTS_NAME) and filename.endswith(".bin"):
         yield filename[:-4] + ".safetensors"
+
+
+def _get_license_from_hf_hub(model_id: str | None, hf_hub_id: str | None) -> str | None:
+    """Retrieve license information for a model from Hugging Face Hub.
+
+    Fetches the license field from the model card metadata on Hugging Face Hub
+    for the specified model. Returns None if the model is not found, if
+    huggingface_hub is not installed, or if the model is marked as "untrained".
+
+    Args:
+        model_id: The model identifier/name. In the case of None we assume an untrained model.
+        hf_hub_id: The Hugging Face Hub organization/user ID. If it is None,
+            we will return None as we cannot infer the license terms.
+
+    Returns:
+        The license string in lowercase if found, None otherwise.
+
+    Note:
+        Requires huggingface_hub package to be installed. Will log a warning
+        and return None if the package is not available.
+    """
+    if not has_hf_hub(True):
+        msg = "For updated license information run `pip install huggingface_hub`."
+        _logger.warning(msg=msg)
+        return None
+
+    if not (model_id and hf_hub_id):
+        return None
+
+    repo_id: str = hf_hub_id + model_id
+
+    try:
+        info = model_info(repo_id=repo_id)
+
+    except RepositoryNotFoundError:
+        # TODO: any wish what happens here? @rwightman
+        print(repo_id)
+        return None
+
+    license = info.card_data.get("license").lower() if info.card_data else None
+    return license

timm/models/cspnet.py

@@ -22,6 +22,7 @@
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from timm.layers import ClassifierHead, ConvNormAct, DropPath, calculate_drop_path_rates, get_attn, create_act_layer, make_divisible
 from ._builder import build_model_with_cfg
+from ._hub import _get_license_from_hf_hub
 from ._manipulate import named_apply, MATCH_PREV_GROUP
 from ._registry import register_model, generate_default_cfgs
 
@@ -1011,82 +1012,97 @@ def _cfg(url='', **kwargs):
         'crop_pct': 0.887, 'interpolation': 'bilinear',
         'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
         'first_conv': 'stem.conv1.conv', 'classifier': 'head.fc',
+        'license': _get_license_from_hf_hub(kwargs.pop('model_id', None), kwargs.get('hf_hub_id')),
         **kwargs
     }
 
-
 default_cfgs = generate_default_cfgs({
     'cspresnet50.ra_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cspresnet50.ra_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/cspresnet50_ra-d3e8d487.pth'),
     'cspresnet50d.untrained': _cfg(),
     'cspresnet50w.untrained': _cfg(),
     'cspresnext50.ra_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cspresnext50.ra_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/cspresnext50_ra_224-648b4713.pth',
     ),
     'cspdarknet53.ra_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cspdarknet53.ra_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/cspdarknet53_ra_256-d05c7c21.pth'),
 
     'darknet17.untrained': _cfg(),
     'darknet21.untrained': _cfg(),
     'sedarknet21.untrained': _cfg(),
     'darknet53.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='darknet53.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/darknet53_256_c2ns-3aeff817.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=1.0),
     'darknetaa53.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='darknetaa53.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/darknetaa53_c2ns-5c28ec8a.pth',
         test_input_size=(3, 288, 288), test_crop_pct=1.0),
 
     'cs3darknet_s.untrained': _cfg(interpolation='bicubic'),
     'cs3darknet_m.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3darknet_m.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3darknet_m_c2ns-43f06604.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=0.95,
     ),
     'cs3darknet_l.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3darknet_l.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3darknet_l_c2ns-16220c5d.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=0.95),
     'cs3darknet_x.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3darknet_x.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3darknet_x_c2ns-4e4490aa.pth',
         interpolation='bicubic', crop_pct=0.95, test_input_size=(3, 288, 288), test_crop_pct=1.0),
 
     'cs3darknet_focus_s.ra4_e3600_r256_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3darknet_focus_s.ra4_e3600_r256_in1k',
         mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
         interpolation='bicubic', test_input_size=(3, 320, 320), test_crop_pct=1.0),
     'cs3darknet_focus_m.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3darknet_focus_m.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3darknet_focus_m_c2ns-e23bed41.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=0.95),
     'cs3darknet_focus_l.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3darknet_focus_l.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3darknet_focus_l_c2ns-65ef8888.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=0.95),
     'cs3darknet_focus_x.untrained': _cfg(interpolation='bicubic'),
 
     'cs3sedarknet_l.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3sedarknet_l.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3sedarknet_l_c2ns-e8d1dc13.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=0.95),
     'cs3sedarknet_x.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3sedarknet_x.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3sedarknet_x_c2ns-b4d0abc0.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=1.0),
 
     'cs3sedarknet_xdw.untrained': _cfg(interpolation='bicubic'),
 
     'cs3edgenet_x.c2_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3edgenet_x.c2_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3edgenet_x_c2-2e1610a9.pth',
         interpolation='bicubic', test_input_size=(3, 288, 288), test_crop_pct=1.0),
     'cs3se_edgenet_x.c2ns_in1k': _cfg(
         hf_hub_id='timm/',
+        model_id='cs3se_edgenet_x.c2ns_in1k',
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tpu-weights/cs3se_edgenet_x_c2ns-76f8e3ac.pth',
         interpolation='bicubic', crop_pct=0.95, test_input_size=(3, 320, 320), test_crop_pct=1.0),
 })
@@ -1204,4 +1220,4 @@ def cs3edgenet_x(pretrained=False, **kwargs) -> CspNet:
 
 @register_model
 def cs3se_edgenet_x(pretrained=False, **kwargs) -> CspNet:
-    return _create_cspnet('cs3se_edgenet_x', pretrained=pretrained, **kwargs)
+    return _create_cspnet('cs3se_edgenet_x', pretrained=pretrained, **kwargs)