SALT-NLP/Design2Code

1---2license: odc-by3---4 5This dataset consists of 484 webpages from the C4 validation set, serving the purpose of testing multimodal LLMs on converting visual designs into code implementations.6 7Each example is a pair of source HTML and screenshot ({id}.html and {id}.png).8See the dataset in the huggingface format [here](https://huggingface.co/datasets/SALT-NLP/Design2Code-hf).9 10Note that all images in these webpages are replaced by a placeholder image (rick.jpg)11 12Please refer to our [project page](https://salt-nlp.github.io/Design2Code/) and [our paper](arxiv.org/abs/2403.03163) for more information.13 14# Example Usage15 16For example, you can generate predictions using [HuggingFaceM4/VLM_WebSight_finetuned](https://huggingface.co/HuggingFaceM4/VLM_WebSight_finetuned).17 18```python19import torch20from PIL import Image21from transformers import AutoModelForCausalLM, AutoProcessor22from transformers.image_utils import to_numpy_array, PILImageResampling, ChannelDimension23from transformers.image_transforms import resize, to_channel_dimension_format24from gpt4v_utils import cleanup_response25from tqdm import tqdm 26import os27 28DEVICE = torch.device("cuda")29HF_TOKEN = "..." # Your HF_TOKEN30 31PROCESSOR = AutoProcessor.from_pretrained(32    "HuggingFaceM4/VLM_WebSight_finetuned",33    token=HF_TOKEN34)35MODEL = AutoModelForCausalLM.from_pretrained(36    "HuggingFaceM4/VLM_WebSight_finetuned",37    token=HF_TOKEN,38    trust_remote_code=True,39    torch_dtype=torch.bfloat16,40).to(DEVICE)41 42print ("parameter count: ", MODEL.num_parameters())43 44image_seq_len = MODEL.config.perceiver_config.resampler_n_latents45BOS_TOKEN = PROCESSOR.tokenizer.bos_token46BAD_WORDS_IDS = PROCESSOR.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids47 48def convert_to_rgb(image):49    # `image.convert("RGB")` would only work for .jpg images, as it creates a wrong background50    # for transparent images. The call to `alpha_composite` handles this case51    if image.mode == "RGB":52        return image53 54    image_rgba = image.convert("RGBA")55    background = Image.new("RGBA", image_rgba.size, (255, 255, 255))56    alpha_composite = Image.alpha_composite(background, image_rgba)57    alpha_composite = alpha_composite.convert("RGB")58    return alpha_composite59 60# The processor is the same as the Idefics processor except for the BILINEAR interpolation,61# so this is a hack in order to redefine ONLY the transform method62def custom_transform(x):63    x = convert_to_rgb(x)64    x = to_numpy_array(x)65    x = resize(x, (960, 960), resample=PILImageResampling.BILINEAR)66    x = PROCESSOR.image_processor.rescale(x, scale=1 / 255)67    x = PROCESSOR.image_processor.normalize(68        x,69        mean=PROCESSOR.image_processor.image_mean,70        std=PROCESSOR.image_processor.image_std71    )72    x = to_channel_dimension_format(x, ChannelDimension.FIRST)73    x = torch.tensor(x)74    return x75 76inputs = PROCESSOR.tokenizer(77    f"{BOS_TOKEN}<fake_token_around_image>{'<image>' * image_seq_len}<fake_token_around_image>",78    return_tensors="pt",79    add_special_tokens=False,80)81 82 83test_data_dir = "/path/to/Design2Code"84predictions_dir = "/path/to/VLM_WebSight_predictions"85 86for filename in tqdm(os.listdir(test_data_dir)):87    if filename.endswith(".png"):88        image_path = os.path.join(test_data_dir, filename)89        with Image.open(image_path) as image:90            inputs["pixel_values"] = PROCESSOR.image_processor([image], transform=custom_transform)91        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}92        generated_ids = MODEL.generate(**inputs, bad_words_ids=BAD_WORDS_IDS, max_length=4096)93        generated_text = PROCESSOR.batch_decode(generated_ids, skip_special_tokens=True)[0]94        generated_text = cleanup_response(generated_text)95 96        with open(os.path.join(predictions_dir, filename.replace(".png", ".html")), "w", encoding='utf-8') as f:97            f.write(generated_text)98```