-
Notifications
You must be signed in to change notification settings - Fork 0
/
app.py
executable file
·135 lines (122 loc) · 4.73 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
# import os
# os.system("pip uninstall -y mmcv-full")
# os.system("mim install 'mmengine>=0.6.0'")
# # os.system("pip install mmcv==2.0.1 -f https://download.openmmlab.com/mmcv/dist/cu118/torch2.0/index.html")
# os.system("mim install 'mmcv-lite==2.0.1'")
# os.system("mim install 'mmdet>=3.0.0,<4.0.0'")
# os.system("mim install 'mmyolo'")
# os.system("pip install -e .")
#
# import argparse
# import os.path as osp
#
# from mmengine.config import Config, DictAction
# from mmengine.runner import Runner
# from mmengine.dataset import Compose
# from mmyolo.registry import RUNNERS
#
# from tools.demo import demo
#
#
# def parse_args():
# parser = argparse.ArgumentParser(
# description='YOLO-World Demo')
# parser.add_argument('--config', default='configs/pretrain/yolo_world_xl_t2i_bn_2e-4_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py')
# parser.add_argument('--checkpoint', default='yolo_world_v2_xl_obj365v1_goldg_cc3mlite_pretrain.pth')
# parser.add_argument(
# '--work-dir',
# help='the directory to save the file containing evaluation metrics')
# parser.add_argument(
# '--cfg-options',
# nargs='+',
# action=DictAction,
# help='override some settings in the used config, the key-value pair '
# 'in xxx=yyy format will be merged into config file. If the value to '
# 'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
# 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
# 'Note that the quotation marks are necessary and that no white space '
# 'is allowed.')
# args = parser.parse_args()
# return args
#
#
# if __name__ == '__main__':
# args = parse_args()
#
# # load config
# cfg = Config.fromfile(args.config)
# if args.cfg_options is not None:
# cfg.merge_from_dict(args.cfg_options)
#
# if args.work_dir is not None:
# cfg.work_dir = args.work_dir
# elif cfg.get('work_dir', None) is None:
# cfg.work_dir = osp.join('./work_dirs',
# osp.splitext(osp.basename(args.config))[0])
#
# cfg.load_from = args.checkpoint
#
# if 'runner_type' not in cfg:
# runner = Runner.from_cfg(cfg)
# else:
# runner = RUNNERS.build(cfg)
#
# runner.call_hook('before_run')
# runner.load_or_resume()
# pipeline = cfg.test_dataloader.dataset.pipeline
# runner.pipeline = Compose(pipeline)
# runner.model.eval()
# demo(runner, args, cfg)
import argparse
import os
import os.path as osp
import cv2
import torch
import numpy as np
from PIL import Image
from mmengine.config import Config, DictAction
from mmengine.runner import Runner
from mmengine.dataset import Compose
from mmyolo.registry import RUNNERS
from tools.demo import run_image # 导入 run_image 函数
def parse_args():
parser = argparse.ArgumentParser(description='YOLO-World Demo')
parser.add_argument('image_path', help='Path to input image')
parser.add_argument('text', help='Comma separated list of classes to detect')
parser.add_argument('--config',
default='configs/pretrain/yolo_world_xl_t2i_bn_2e-4_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py')
parser.add_argument('--checkpoint', default='yolo_world_v2_xl_obj365v1_goldg_cc3mlite_pretrain.pth')
parser.add_argument('--output', default='./output.png', help='Path to save the output image')
parser.add_argument('--work-dir', help='Directory to save the evaluation metrics')
parser.add_argument('--cfg-options', nargs='+', action=DictAction, help='Override settings in the config')
return parser.parse_args()
def main():
args = parse_args()
# Load configuration
cfg = Config.fromfile(args.config)
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
if args.work_dir is not None:
cfg.work_dir = args.work_dir
elif cfg.get('work_dir', None) is None:
cfg.work_dir = osp.join('./work_dirs', osp.splitext(osp.basename(args.config))[0])
cfg.load_from = args.checkpoint
# Initialize the model
if 'runner_type' not in cfg:
runner = Runner.from_cfg(cfg)
else:
runner = RUNNERS.build(cfg)
runner.call_hook('before_run')
runner.load_or_resume()
pipeline = cfg.test_dataloader.dataset.pipeline
runner.pipeline = Compose(pipeline)
runner.model.eval()
# Run image detection
image = Image.open(args.image_path) # Open image from file path
output_image = run_image(runner, image, args.text, max_num_boxes=100, score_thr=0.05, nms_thr=0.5,
image_path=args.output)
# Save the output image
output_image.save(args.output)
print(f"Output saved to {args.output}")
if __name__ == '__main__':
main()