dragonSwing/annotate-anything

1import os
2
3# Configurations
4tag2text_dict = {
5    "swin_14m": {
6        "checkpoint_url": "https://huggingface.co/spaces/xinyu1205/Tag2Text/resolve/main/tag2text_swin_14m.pth",
7        "checkpoint_file": "tag2text_swin_14m.pth",
8    }
9}
10
11sam_dict = {
12    "default": {
13        "checkpoint_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
14        "checkpoint_file": "sam_vit_h_4b8939.pth",
15    },
16    "vit_h": {
17        "checkpoint_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
18        "checkpoint_file": "sam_vit_h_4b8939.pth",
19    },
20    "vit_l": {
21        "checkpoint_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
22        "checkpoint_file": "sam_vit_l_0b3195.pth",
23    },
24    "vit_b": {
25        "checkpoint_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
26        "checkpoint_file": "sam_vit_b_01ec64.pth",
27    },
28}
29
30dino_dict = {
31    "swinb": {
32        "repo_id": "ShilongLiu/GroundingDINO",
33        "config_file": "GroundingDINO_SwinB.cfg.py",
34        "checkpoint_file": "groundingdino_swinb_cogcoor.pth",
35    },
36    "swint_ogc": {
37        "repo_id": "ShilongLiu/GroundingDINO",
38        "config_file": "GroundingDINO_SwinT_OGC.cfg.py",
39        "checkpoint_file": "groundingdino_swint_ogc.pth",
40    },
41}
42
43default_sam = "default"
44default_tag2text = "swin_14m"
45default_dino = "swint_ogc"
46
47root_dir = os.path.dirname(os.path.abspath(__file__))
48weight_dir = "weights"
49abs_weight_dir = os.path.join(root_dir, weight_dir)
50
51tag2text_checkpoint = "tag2text_swin_14m.pth"
52tag2text_url = "https://huggingface.co/spaces/xinyu1205/Tag2Text/resolve/main/tag2text_swin_14m.pth"
53sam_checkpoint = "sam_vit_h_4b8939.pth"
54sam_url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
55output_dir = "outputs"
56
57dino_config_file = "GroundingDINO_SwinB.cfg.py"
58dino_repo_id = "ShilongLiu/GroundingDINO"
59dino_checkpoint = "groundingdino_swinb_cogcoor.pth"
60
61iou_threshold = 0.5
62box_threshold = 0.3
63text_threshold = 0.25
64
65# filter out attributes and action categories which are difficult to grounding
66delete_tag_index = []
67for i in range(3012, 3429):
68    delete_tag_index.append(i)
69