129 lines
2.3 KiB
Python
129 lines
2.3 KiB
Python
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()
|