Seeking Advice on Vision Model Optimization for an Agent That Controls a Computer

Hello everyone,

I am currently developing a chatbot that can connect to cloud APIs, CLI tools, or local models.

One of its small features is the ability to analyze the user’s mouse movement coordinates, click positions, and button or icon images. The goal is to let the chatbot learn how to help operate a computer screen.

At the moment, I am using YOLOX (Apache 2.0) as the base model. I annotate new training data myself and fine-tune the model so it can understand what UI buttons are.

However, I have run into a problem. Web buttons are usually wide and often contain text, while desktop application buttons are often icon-only. When I created my initial annotations, I did not pay enough attention to this difference. As a result, the fine-tuned model has become somewhat biased: it recognizes text-based buttons more easily, but struggles with icon-only buttons.

I have tried using OpenCV together with confidence scores to compensate for this, but it only helps slightly. The effect is still limited. Because of that, I am considering rebuilding or expanding the training dataset, especially since I would like to keep the project compatible with the Apache 2.0 license.

I have searched for information and also asked AI tools, but I received many different suggestions. Since I am an individual developer, my time is limited, and model training also takes time and resources. So I would like to ask for advice from people with more experience:

When expanding the training dataset, would you recommend adding a large amount of icon-only button data?

Or should I continue training with a balanced 5:5 ratio between text-based buttons and icon-only buttons?

Or is there a better annotation strategy or training approach for this kind of UI control detection task?

I am currently training the model on an NVIDIA GeForce RTX 4050.

Below is an example of the operation plan generated by the system. These actions are stored in the bot’s memory:

A safe plan has been generated.
After confirmation, the system will first try to operate elements inside the current window.
Run: learn-1780353528723374400, 5 steps.

1. click native-window, screen coordinates (1151, 1050), window: unknown, process: explorer.exe
2. click Google Chrome, screen coordinates (974, 611), window: Google Chrome, process: chrome.exe
3. click New Tab - Google Chrome, screen coordinates (425, 598), window: New Tab - Google Chrome, process: chrome.exe
4. click ChatGPT - GitHub - Google Chrome, screen coordinates (210, 202), window: ChatGPT - GitHub - Google Chrome, process: chrome.exe
5. click ChatGPT - Google Chrome, screen coordinates (1549, 29), window: ChatGPT - Google Chrome, process: chrome.exe

During execution, elements inside the current window are handled through selectors first.
External browser windows use native screen coordinates.
If no matching element can be found inside the current window, the system falls back to the demonstrated coordinates.

If reviewing the code would help provide better advice, the GitHub repository is here:

Thank you in advance for any suggestions or guidance.