Agent-S grounding.py

myeongjaechoi 2025. 3. 30. 16:34

좌표 생성 메서드 (generate_coords)

def generate_coords(self, ref_expr: str, obs: Dict) -> List[int]:
    self.grounding_model.reset()
    prompt = f"Query:{ref_expr}\nOutput only the coordinate of one point in your response.\n"
    self.grounding_model.add_message(
        text_content=prompt, image_content=obs["screenshot"], put_text_last=True
    )
    response = call_llm_safe(self.grounding_model)
    print("RAW GROUNDING MODEL RESPONSE:", response)
    numericals = re.findall(r"\d+", response)
    assert len(numericals) >= 2
    return [int(numericals[0]), int(numericals[1])]

ref_expr: 찾고자 하는 요소에 대한 설명(예: "로그인 버튼")
obs: 화면 스크린샷을 포함한 관찰 데이터
언어 모델이 설명에 해당하는 화면 요소의 좌표를 생성하고 반환합니다

OCR 요소 추출 메서드 (get_ocr_elements)

이 메서드는 이미지에서 텍스트를 인식하고 각 텍스트 요소의 위치 정보를 추출합니다:

b64_image_data: 화면 스크린샷 이미지 데이터
OCR(광학 문자 인식)을 사용하여 이미지에서 텍스트를 추출하고, 각 텍스트의 위치와 크기 정보를 포함한 목록을 반환합니다

def get_ocr_elements(self, b64_image_data: str) -> Tuple[str, List]:
    image = Image.open(BytesIO(b64_image_data))
    image_data = pytesseract.image_to_data(image, output_type=Output.DICT)
    
    # 텍스트 정리 (앞뒤 공백과 특수문자 제거)
    for i, word in enumerate(image_data["text"]):
        image_data["text"][i] = re.sub(
            r"^[^a-zA-Z\s.,!?;:\-\+]+|[^a-zA-Z\s.,!?;:\-\+]+$", "", word
        )
    
    # OCR 요소 생성 및 테이블 형태로 정리
    ocr_elements = []
    ocr_table = "Text Table:\nWord id\tText\n"
    grouping_map = defaultdict(list)
    ocr_id = 0
    
    for i in range(len(image_data["text"])):
        block_num = image_data["block_num"][i]
        if image_data["text"][i]:
            grouping_map[block_num].append(image_data["text"][i])
            ocr_table += f"{ocr_id}\t{image_data['text'][i]}\n"
            ocr_elements.append({
                "id": ocr_id,
                "text": image_data["text"][i],
                "group_num": block_num,
                "word_num": len(grouping_map[block_num]),
                "left": image_data["left"][i],
                "top": image_data["top"][i],
                "width": image_data["width"][i],
                "height": image_data["height"][i],
            })
            ocr_id += 1
    
    return ocr_table, ocr_elements

텍스트 좌표 생성 메서드 (generate_text_coords)

이 메서드는 특정 텍스트 문구의 좌표를 찾습니다

phrase: 찾을 텍스트 문구
obs: 화면 스크린샷을 포함한 관찰 데이터
alignment: 좌표의 정렬 방식 ("start"는 문구의 시작, "end"는 문구의 끝)
OCR로 추출된 텍스트 중에서 주어진 문구와 일치하는 텍스트의 좌표를 반환합니다

def generate_text_coords(self, phrase: str, obs: Dict, alignment: str = "") -> List[int]:
    ocr_table, ocr_elements = self.get_ocr_elements(obs["screenshot"])
    
    alignment_prompt = ""
    if alignment == "start":
        alignment_prompt = "**Important**: Output the word id of the FIRST word in the provided phrase.\n"
    elif alignment == "end":
        alignment_prompt = "**Important**: Output the word id of the LAST word in the provided phrase.\n"
    
    # 언어 모델 프롬프트 설정
    self.text_span_agent.reset()
    self.text_span_agent.add_message(
        alignment_prompt + "Phrase: " + phrase + "\n" + ocr_table, role="user"
    )
    self.text_span_agent.add_message(
        "Screenshot:\n", image_content=obs["screenshot"], role="user"
    )
    
    # 응답 처리
    response = call_llm_safe(self.text_span_agent)
    print("TEXT SPAN AGENT RESPONSE:", response)
    numericals = re.findall(r"\d+", response)
    if len(numericals) > 0:
        text_id = int(numericals[-1])
    else:
        text_id = 0
    elem = ocr_elements[text_id]
    
    # 좌표 계산
    if alignment == "start":
        coords = [elem["left"], elem["top"] + (elem["height"] // 2)]
    elif alignment == "end":
        coords = [elem["left"] + elem["width"], elem["top"] + (elem["height"] // 2)]
    else:
        coords = [
            elem["left"] + (elem["width"] // 2),
            elem["top"] + (elem["height"] // 2),
        ]
    return coords

좌표 할당 메서드 (assign_coordinates)

plan: 실행할 액션 계획
obs: 화면 스크린샷을 포함한 관찰 데이터
액션 유형에 따라 적절한 좌표 생성 메서드를 호출하여 좌표를 할당합니다

def assign_coordinates(self, plan: str, obs: Dict):
    # 이전 좌표 초기화
    self.coords1, self.coords2 = None, None
    
    try:
        # 함수 이름과 인자 추출
        action = parse_single_code_from_string(plan.split("Grounded Action")[-1])
        function_name = re.match(r"(\w+\.\w+)\(", action).group(1)
        args = self.parse_function_args(action)
    except Exception as e:
        raise RuntimeError(f"Error in parsing grounded action: {e}") from e
    
    # 액션 유형에 따라 좌표 생성
    if (
        function_name in ["agent.click", "agent.type", "agent.scroll"]
        and len(args) >= 1
        and args[0] != None
    ):
        self.coords1 = self.generate_coords(args[0], obs)
    elif function_name == "agent.drag_and_drop" and len(args) >= 2:
        self.coords1 = self.generate_coords(args[0], obs)
        self.coords2 = self.generate_coords(args[1], obs)
    elif function_name == "agent.highlight_text_span" and len(args) >= 2:
        self.coords1 = self.generate_text_coords(args[0], obs, alignment="start")
        self.coords2 = self.generate_text_coords(args[1], obs, alignment="end")

클릭 액션 (click)

element_description: 클릭할 요소에 대한 설명
num_clicks: 클릭 횟수
button_type: 마우스 버튼 유형
hold_keys: 클릭하는 동안 누르고 있을 키 목록
PyAutoGUI 라이브러리를 사용하여 클릭 명령을 생성합니다

@agent_action
def click(
    self,
    element_description: str,
    num_clicks: int = 1,
    button_type: str = "left",
    hold_keys: List = [],
):
    """요소 클릭
    인자:
        element_description:str, 클릭할 요소에 대한 상세한 설명. 최소 완전한 문장이어야 함
        num_clicks:int, 요소를 클릭할 횟수
        button_type:str, 사용할 마우스 버튼 ("left", "middle", "right" 중 하나)
        hold_keys:List, 클릭하는 동안 누르고 있을 키 목록
    """
    x, y = self.resize_coordinates(self.coords1)
    command = "import pyautogui; "
    
    # 키 누르기
    for k in hold_keys:
        command += f"pyautogui.keyDown({repr(k)}); "
    
    # 클릭 명령 추가
    command += f"""import pyautogui; pyautogui.click({x}, {y}, clicks={num_clicks}, button={repr(button_type)}); """
    
    # 키 떼기
    for k in hold_keys:
        command += f"pyautogui.keyUp({repr(k)}); "
    
    return command

애플리케이션 전환 액션 (switch_applications)

app_code: 전환할 애플리케이션의 이름
각 운영체제(Mac, Ubuntu, Windows)에 맞는 전환 명령을 생성합니다

@agent_action
def switch_applications(self, app_code):
    """다른 열린 애플리케이션으로 전환
    인자:
        app_code:str 전환할 애플리케이션의 코드명
    """
    if self.platform == "mac":
        return f"import pyautogui; import time; pyautogui.hotkey('command', 'space', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"
    elif self.platform == "ubuntu":
        return UBUNTU_APP_SETUP.replace("APP_NAME", app_code)
    elif self.platform == "windows":
        return f"import pyautogui; import time; pyautogui.hotkey('win', 'd', interval=0.5); pyautogui.typewrite({repr(app_code)}); pyautogui.press('enter'); time.sleep(1.0)"

텍스트 입력 액션 (type)

element_description: 텍스트를 입력할 요소에 대한 설명
text: 입력할 텍스트
overwrite: 기존 텍스트를 덮어쓸지 여부
enter: 텍스트 입력 후 엔터키를 누를지 여부
요소 위치에 클릭한 후 텍스트를 입력하는 명령을 생성합니다13

@agent_action
def type(
    self,
    element_description: Optional[str] = None,
    text: str = "",
    overwrite: bool = False,
    enter: bool = False,
):
    """특정 요소에 텍스트 입력
    인자:
        element_description:str, 텍스트를 입력할 요소에 대한 상세한 설명
        text:str, 입력할 텍스트
        overwrite:bool, 기존 텍스트를 덮어쓸지 여부
        enter:bool, 텍스트 입력 후 엔터키를 누를지 여부
    """
    if self.coords1 is not None:
        x, y = self.resize_coordinates(self.coords1)
        
        command = "import pyautogui; "
        command += f"pyautogui.click({x}, {y}); "
        
        if overwrite:
            command += f"pyautogui.hotkey('ctrl', 'a'); pyautogui.press('backspace'); "
            
        command += f"pyautogui.write({repr(text)}); "
        
        if enter:
            command += "pyautogui.press('enter'); "
    else:
        command = "import pyautogui; "
        
        if overwrite:
            command += f"pyautogui.hotkey('ctrl', 'a'); pyautogui.press('backspace'); "
            
        command += f"pyautogui.write({repr(text)}); "
        
        if enter:
            command += "pyautogui.press('enter'); "
            
    return command