cockpit-ui-grounding/scripts/run_grounding.py

129 lines
2.3 KiB
Python
Raw Normal View History

2026-08-24 16:29:35 +08:00
import argparse
import json
from cockpit_grounding.models.factory import create_grounder
from cockpit_grounding.grounding.predictor import (
build_grounding_prompt,
parse_grounding_output,
)
from cockpit_grounding.vision.visualize import (
visualize_grounding,
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--image",
required=True,
)
parser.add_argument(
"--target",
required=True,
)
parser.add_argument(
"--output",
required=True,
)
parser.add_argument(
"--model",
required=True,
help="Local model path",
)
parser.add_argument(
"--backend",
choices=("qwen3vl", "qwen35"),
default="qwen3vl",
help="Model backend (default: qwen3vl)",
)
args = parser.parse_args()
# ----------------------------
# Load model
# ----------------------------
grounder = create_grounder(
backend=args.backend,
model_path=args.model,
)
# ----------------------------
# Prompt
# ----------------------------
prompt = build_grounding_prompt(
args.target
)
# ----------------------------
# Inference
# ----------------------------
raw_output = grounder.generate(
args.image,
prompt,
)
print()
print("========== RAW MODEL OUTPUT ==========")
print(raw_output)
# ----------------------------
# Parse bbox
# ----------------------------
result = parse_grounding_output(
raw_output
)
# ----------------------------
# Convert + draw
# ----------------------------
pixel_result = visualize_grounding(
args.image,
result,
args.output,
)
final_result = {
"target": args.target,
"bbox_relative": [
result.x1,
result.y1,
result.x2,
result.y2,
],
**pixel_result,
}
print()
print("========== FINAL RESULT ==========")
print(
json.dumps(
final_result,
indent=2,
ensure_ascii=False,
)
)
print()
print(
f"Result image: {args.output}"
)
if __name__ == "__main__":
main()