| from .clip_models import CLIPModel |
| from .imagenet_models import ImagenetModel |
|
|
|
|
| VALID_NAMES = [ |
| 'Imagenet:resnet18', |
| 'Imagenet:resnet34', |
| 'Imagenet:resnet50', |
| 'Imagenet:resnet101', |
| 'Imagenet:resnet152', |
| 'Imagenet:vgg11', |
| 'Imagenet:vgg19', |
| 'Imagenet:swin-b', |
| 'Imagenet:swin-s', |
| 'Imagenet:swin-t', |
| 'Imagenet:vit_b_16', |
| 'Imagenet:vit_b_32', |
| 'Imagenet:vit_l_16', |
| 'Imagenet:vit_l_32', |
|
|
| 'CLIP:RN50', |
| 'CLIP:RN101', |
| 'CLIP:RN50x4', |
| 'CLIP:RN50x16', |
| 'CLIP:RN50x64', |
| 'CLIP:ViT-B/32', |
| 'CLIP:ViT-B/16', |
| 'CLIP:ViT-L/14', |
| 'CLIP:ViT-L/14@336px', |
| ] |
|
|
|
|
|
|
|
|
|
|
| def get_model(name): |
| assert name in VALID_NAMES |
| if name.startswith("Imagenet:"): |
| return ImagenetModel(name[9:]) |
| elif name.startswith("CLIP:"): |
| return CLIPModel(name[5:]) |
| else: |
| assert False |
|
|
