If I understood you right, your code works, but you need to get the TokenAccount address from the wallet right?
What you can do is: call getTokenAccountsByOwner(pubKey) and then check all the token accounts TokenAccountInfo mint if its the mint of your specific token.
Then you can call getTokenAccountBalance on that account like you already do.
To make it more performant you can directly specify the mint in the getTokenAccountsByOwner function. Then you can save the token account in the client and only call getTokenAccountBalance form then on.