| """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) |
| |
| 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 |
|
|
