Spaces:
Running
Running
| import tome | |
| import timm | |
| import gradio as gr | |
| from PIL import Image | |
| from torchvision import transforms | |
| from torchvision.transforms.functional import InterpolationMode | |
| model_name = "vit_large_patch16_384" | |
| print("Started Downloading:", model_name) | |
| model = timm.create_model(model_name, pretrained=True) | |
| print("Finished Downloading:", model_name) | |
| tome.patch.timm(model, trace_source=True) | |
| input_size = model.default_cfg["input_size"][1] | |
| # Make sure the transform is correct for your model! | |
| transform_list = [ | |
| transforms.Resize(int((256 / 224) * input_size), interpolation=InterpolationMode.BICUBIC), | |
| transforms.CenterCrop(input_size) | |
| ] | |
| # The visualization and model need different transforms | |
| transform_vis = transforms.Compose(transform_list) | |
| transform_norm = transforms.Compose(transform_list + [ | |
| transforms.ToTensor(), | |
| transforms.Normalize(model.default_cfg["mean"], model.default_cfg["std"]), | |
| ]) | |
| def process_image(img, r=25, layers=1): | |
| img = Image.fromarray(img.astype('uint8'), 'RGB') | |
| img_vis = transform_vis(img) | |
| img_norm = transform_norm(img) | |
| # from the paper: | |
| # r can take the following forms: | |
| # - int: A constant number of tokens per layer. | |
| # - Tuple[int, float]: A pair of r, inflection. | |
| # Inflection describes there the the reduction / layer should trend | |
| # upward (+1), downward (-1), or stay constant (0). A value of (r, 0) | |
| # is as providing a constant r. (r, -1) is what we describe in the paper | |
| # as "decreasing schedule". Any value between -1 and +1 is accepted. | |
| # - List[int]: A specific number of tokens per layer. For extreme granularity. | |
| if layers != 1: | |
| r = [r] * layers | |
| print(r) | |
| model.r = r | |
| _ = model(img_norm[None, ...]) | |
| source = model._tome_info["source"] | |
| # print(f"{source.shape[1]} tokens at the end") | |
| return tome.make_visualization(img_vis, source, patch_size=16, class_token=True) | |
| iface = gr.Interface( | |
| fn=process_image, | |
| inputs=[ | |
| "image", | |
| gr.inputs.Slider(0, 50, step=1, label="r value (the amount of reduction. See paper for details.)"), | |
| gr.inputs.Slider(1, 50, step=1, label="layers (1 means r is applied to all layers)"), | |
| ], | |
| outputs="image", | |
| examples=[ | |
| ["images/husky.png", 25, 1], | |
| ["images/husky.png", 25, 8], | |
| ["images/husky.png", 25, 16], | |
| ["images/husky.png", 25, 22], | |
| ] | |
| ) | |
| iface.launch() | |