File size: 2,768 Bytes
ed79a7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """Opt-in experimental icon labels; labels are evidence, never click commands."""
import io,json
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import load_file
from transformers.models.idefics3.configuration_idefics3 import Idefics3VisionConfig
from transformers.models.idefics3.modeling_idefics3 import Idefics3VisionTransformer
from miniswardbower.browser.dom_distiller import DOMDistiller
class VisualIconDistiller(DOMDistiller):
def __init__(self, directory):
super().__init__(draw_overlays=False)
directory=Path(directory)
self.labels=json.loads((directory/'report.json').read_text())['labels']
config=Idefics3VisionConfig.from_json_file(directory/'vision_config.json');config._attn_implementation='sdpa'
self.vision=Idefics3VisionTransformer(config)
state=load_file(str(directory/'icon_specialist.safetensors'))
self.vision.load_state_dict({k[7:]:v for k,v in state.items() if k.startswith('vision.')},strict=True)
self.head=torch.nn.Linear(config.hidden_size,len(self.labels))
self.head.load_state_dict({k[5:]:v for k,v in state.items() if k.startswith('head.')},strict=True)
self.vision=self.vision.cuda().bfloat16().eval();self.head=self.head.cuda().eval()
self.observations=[]
@torch.inference_mode()
def classify(self, image):
a=torch.from_numpy(np.asarray(image.convert('RGB').resize((512,512))).copy()).permute(2,0,1).float()/255
features=self.vision(pixel_values=((a-.5)/.5)[None].cuda().bfloat16()).last_hidden_state.mean(1).float()
probs=self.head(F.normalize(features,dim=-1)).softmax(-1)[0]
return self.labels[int(probs.argmax())],float(probs.max())
async def extract_pruned_tree(self,page):
tree=await super().extract_pruned_tree(page)
# A locator screenshot scrolls offscreen elements into view. Perception must not act.
viewport=Image.open(io.BytesIO(await page.screenshot(full_page=False,scale='css')))
count=0
for elem in tree.elements:
x,y,w,h=elem.bbox
if elem.aria_label or elem.tag not in ('button','a') or not (0<w<=160 and 0<h<=100):continue
if elem.text and any(c.isalpha() for c in elem.text):continue
if count>=20:break
if x<0 or y<0 or x+w>viewport.width or y+h>viewport.height:continue
name,confidence=self.classify(viewport.crop((x,y,x+w,y+h)));count+=1
self.observations.append(dict(mark=elem.id,label=name,confidence=confidence))
if name!='other':
elem.text=(elem.text or '')+f' [visual guess: {name}; confidence {confidence:.2f}]'
return tree
|