|
| 1 | +"""VLM tab: describe a UI element in words and have a model find it.""" |
| 2 | +from typing import Optional |
| 3 | + |
| 4 | +from PySide6.QtWidgets import ( |
| 5 | + QHBoxLayout, QLabel, QLineEdit, QMessageBox, QPushButton, |
| 6 | + QVBoxLayout, QWidget, |
| 7 | +) |
| 8 | + |
| 9 | +from je_auto_control.gui._i18n_helpers import TranslatableMixin |
| 10 | +from je_auto_control.gui.language_wrapper.multi_language_wrapper import ( |
| 11 | + language_wrapper, |
| 12 | +) |
| 13 | +from je_auto_control.utils.vision.backends.base import VLMNotAvailableError |
| 14 | +from je_auto_control.utils.vision.vlm_api import ( |
| 15 | + click_by_description, locate_by_description, |
| 16 | +) |
| 17 | + |
| 18 | + |
| 19 | +def _t(key: str) -> str: |
| 20 | + return language_wrapper.translate(key, key) |
| 21 | + |
| 22 | + |
| 23 | +class VLMTab(TranslatableMixin, QWidget): |
| 24 | + """Drive a vision-language model to locate or click described elements.""" |
| 25 | + |
| 26 | + def __init__(self, parent: Optional[QWidget] = None) -> None: |
| 27 | + super().__init__(parent) |
| 28 | + self._tr_init() |
| 29 | + self._description = QLineEdit() |
| 30 | + self._model = QLineEdit() |
| 31 | + self._status = QLabel() |
| 32 | + self._last_result = QLabel() |
| 33 | + self._apply_placeholders() |
| 34 | + self._build_layout() |
| 35 | + |
| 36 | + def retranslate(self) -> None: |
| 37 | + TranslatableMixin.retranslate(self) |
| 38 | + self._apply_placeholders() |
| 39 | + |
| 40 | + def _apply_placeholders(self) -> None: |
| 41 | + self._description.setPlaceholderText(_t("vlm_desc_placeholder")) |
| 42 | + self._model.setPlaceholderText(_t("vlm_model_placeholder")) |
| 43 | + |
| 44 | + def _build_layout(self) -> None: |
| 45 | + root = QVBoxLayout(self) |
| 46 | + desc_row = QHBoxLayout() |
| 47 | + desc_row.addWidget(self._tr(QLabel(), "vlm_desc_label")) |
| 48 | + desc_row.addWidget(self._description, stretch=1) |
| 49 | + root.addLayout(desc_row) |
| 50 | + model_row = QHBoxLayout() |
| 51 | + model_row.addWidget(self._tr(QLabel(), "vlm_model_label")) |
| 52 | + model_row.addWidget(self._model, stretch=1) |
| 53 | + root.addLayout(model_row) |
| 54 | + btn_row = QHBoxLayout() |
| 55 | + locate_btn = self._tr(QPushButton(), "vlm_locate") |
| 56 | + locate_btn.clicked.connect(self._on_locate) |
| 57 | + click_btn = self._tr(QPushButton(), "vlm_click") |
| 58 | + click_btn.clicked.connect(self._on_click) |
| 59 | + btn_row.addWidget(locate_btn) |
| 60 | + btn_row.addWidget(click_btn) |
| 61 | + btn_row.addStretch() |
| 62 | + root.addLayout(btn_row) |
| 63 | + root.addWidget(self._last_result) |
| 64 | + root.addWidget(self._status) |
| 65 | + root.addStretch() |
| 66 | + |
| 67 | + def _collect_inputs(self): |
| 68 | + description = self._description.text().strip() |
| 69 | + if not description: |
| 70 | + self._status.setText(_t("vlm_desc_required")) |
| 71 | + return None |
| 72 | + model = self._model.text().strip() or None |
| 73 | + return description, model |
| 74 | + |
| 75 | + def _on_locate(self) -> None: |
| 76 | + inputs = self._collect_inputs() |
| 77 | + if inputs is None: |
| 78 | + return |
| 79 | + description, model = inputs |
| 80 | + try: |
| 81 | + coords = locate_by_description(description, model=model) |
| 82 | + except VLMNotAvailableError as error: |
| 83 | + QMessageBox.warning(self, _t("vlm_locate"), str(error)) |
| 84 | + return |
| 85 | + except (OSError, ValueError, RuntimeError) as error: |
| 86 | + self._status.setText(f"{_t('vlm_error')}: {error}") |
| 87 | + return |
| 88 | + if coords is None: |
| 89 | + self._status.setText(_t("vlm_not_found")) |
| 90 | + self._last_result.setText("") |
| 91 | + return |
| 92 | + self._last_result.setText( |
| 93 | + _t("vlm_result").replace("{x}", str(coords[0])) |
| 94 | + .replace("{y}", str(coords[1])), |
| 95 | + ) |
| 96 | + self._status.setText(_t("vlm_ok")) |
| 97 | + |
| 98 | + def _on_click(self) -> None: |
| 99 | + inputs = self._collect_inputs() |
| 100 | + if inputs is None: |
| 101 | + return |
| 102 | + description, model = inputs |
| 103 | + try: |
| 104 | + ok = click_by_description(description, model=model) |
| 105 | + except VLMNotAvailableError as error: |
| 106 | + QMessageBox.warning(self, _t("vlm_click"), str(error)) |
| 107 | + return |
| 108 | + except (OSError, ValueError, RuntimeError) as error: |
| 109 | + self._status.setText(f"{_t('vlm_error')}: {error}") |
| 110 | + return |
| 111 | + self._status.setText(_t("vlm_ok") if ok else _t("vlm_not_found")) |
0 commit comments