150 lines
5.6 KiB
Python
150 lines
5.6 KiB
Python
"""Test OCR grounding: take a screenshot and find text elements.
|
|
|
|
Usage:
|
|
# Find a specific text on current screen
|
|
python scripts/test_ocr_grounding.py "微信"
|
|
|
|
# Detect ALL text on screen (debug mode)
|
|
python scripts/test_ocr_grounding.py --all
|
|
|
|
# Use a saved screenshot instead of live ADB capture
|
|
python scripts/test_ocr_grounding.py "发送" --image data/screenshots/test.png
|
|
|
|
# Try different engines
|
|
python scripts/test_ocr_grounding.py "微信" --engine easyocr
|
|
python scripts/test_ocr_grounding.py "微信" --engine pytesseract
|
|
|
|
# Also try uiautomator dump (hybrid mode)
|
|
python scripts/test_ocr_grounding.py "微信" --hybrid
|
|
|
|
# Save annotated screenshot with bounding boxes drawn
|
|
python scripts/test_ocr_grounding.py --all --annotate
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import argparse
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from src.grounding.ocr_grounding import OCRGrounding
|
|
|
|
|
|
def annotate_image(img: Image.Image, boxes, query: str = "") -> Image.Image:
|
|
"""Draw bounding boxes on the image for visualization."""
|
|
annotated = img.copy()
|
|
draw = ImageDraw.Draw(annotated)
|
|
|
|
for box in boxes:
|
|
is_match = box.contains_text(query) if query else False
|
|
color = "red" if is_match else "lime"
|
|
width = 3 if is_match else 1
|
|
|
|
draw.rectangle(
|
|
[box.x, box.y, box.x + box.w, box.y + box.h],
|
|
outline=color, width=width,
|
|
)
|
|
label = f"{box.text} ({box.confidence:.0%})"
|
|
draw.text((box.x, box.y - 14), label, fill=color)
|
|
|
|
return annotated
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Test OCR grounding on phone screen")
|
|
parser.add_argument("query", nargs="?", default=None, help="Text to find on screen")
|
|
parser.add_argument("--all", action="store_true", help="Detect all text on screen")
|
|
parser.add_argument("--image", type=str, help="Use saved screenshot instead of ADB")
|
|
parser.add_argument("--engine", type=str, default="auto",
|
|
choices=["auto", "pytesseract", "easyocr"],
|
|
help="OCR engine to use")
|
|
parser.add_argument("--hybrid", action="store_true",
|
|
help="Try uiautomator + OCR hybrid approach")
|
|
parser.add_argument("--annotate", action="store_true",
|
|
help="Save annotated screenshot with bounding boxes")
|
|
args = parser.parse_args()
|
|
|
|
if not args.query and not args.all:
|
|
parser.error("Provide a search query or --all")
|
|
|
|
# Get screenshot
|
|
if args.image:
|
|
print(f"Loading image: {args.image}")
|
|
img = Image.open(args.image)
|
|
else:
|
|
from src.capture import ADBCapture
|
|
cap = ADBCapture()
|
|
info = cap.check_device()
|
|
if not info["connected"]:
|
|
print(f"[FAIL] {info['error']}")
|
|
sys.exit(1)
|
|
print(f"Device: {info['model']} ({info['resolution']})")
|
|
print("Taking screenshot...")
|
|
img = cap.screenshot(save=True)
|
|
|
|
print(f"Image size: {img.width}x{img.height}")
|
|
grounding = OCRGrounding(engine=args.engine)
|
|
|
|
if args.all:
|
|
print(f"\n--- Detecting ALL text (engine={args.engine}) ---\n")
|
|
boxes = grounding.detect_all(img)
|
|
if not boxes:
|
|
print("[WARN] No text detected!")
|
|
else:
|
|
print(f"Found {len(boxes)} text regions:\n")
|
|
for i, box in enumerate(boxes, 1):
|
|
nx, ny = box.center_normalized(img.width, img.height)
|
|
print(f" {i:3d}. '{box.text}'")
|
|
print(f" pixel=({box.cx}, {box.cy}) "
|
|
f"norm=({nx:.3f}, {ny:.3f}) "
|
|
f"conf={box.confidence:.0%}")
|
|
|
|
if args.annotate and boxes:
|
|
out_path = "data/screenshots/annotated_all.png"
|
|
annotated = annotate_image(img, boxes, query=args.query or "")
|
|
annotated.save(out_path)
|
|
print(f"\nAnnotated image saved: {out_path}")
|
|
|
|
if args.query:
|
|
print(f"\n--- Searching for: '{args.query}' (engine={args.engine}) ---\n")
|
|
|
|
if args.hybrid:
|
|
result = grounding.find_text_hybrid(img, args.query)
|
|
else:
|
|
result = grounding.find_text(img, args.query)
|
|
|
|
if result is None:
|
|
print(f"[NOT FOUND] '{args.query}' was not found on screen.")
|
|
print("\nTip: Run with --all to see all detected text.")
|
|
sys.exit(1)
|
|
else:
|
|
nx, ny = result.center_normalized(img.width, img.height)
|
|
print(f"[FOUND] '{result.text}'")
|
|
print(f" Pixel center: ({result.cx}, {result.cy})")
|
|
print(f" Normalized center: ({nx:.4f}, {ny:.4f})")
|
|
print(f" Bounding box: x={result.x} y={result.y} "
|
|
f"w={result.w} h={result.h}")
|
|
print(f" Confidence: {result.confidence:.0%}")
|
|
print()
|
|
print(f" To tap this element:")
|
|
print(f" adb shell input tap {result.cx} {result.cy}")
|
|
|
|
# Show all matches
|
|
all_matches = grounding.find_all_matches(img, args.query)
|
|
if len(all_matches) > 1:
|
|
print(f"\n ({len(all_matches)} total matches found)")
|
|
for i, m in enumerate(all_matches):
|
|
print(f" {i+1}. '{m.text}' at ({m.cx},{m.cy}) conf={m.confidence:.0%}")
|
|
|
|
if args.annotate:
|
|
boxes = grounding.detect_all(img)
|
|
out_path = "data/screenshots/annotated_search.png"
|
|
annotated = annotate_image(img, boxes, query=args.query)
|
|
annotated.save(out_path)
|
|
print(f"\nAnnotated image saved: {out_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|