import os import numpy as np import time import random import torch import torchvision.transforms as transforms import gradio as gr import matplotlib.pyplot as plt from models import get_model from dotmap import DotMap from PIL import Image #os.environ['TERM'] = 'linux' #os.environ['TERMINFO'] = '/etc/terminfo' # args args = DotMap() args.deploy = 'vanilla' args.arch = 'dino_small_patch16' args.no_pretrain = True args.resume = 'https://huggingface.co/hushell/pmf_dinosmall_lr1e-4/resolve/main/best_converted.pth' args.api_key = 'AIzaSyAFkOGnXhy-2ZB0imDvNNqf2rHb98vR_qY' args.cx = '06d75168141bc47f1' # model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = get_model(args) model.to(device) checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu') model.load_state_dict(checkpoint['model'], strict=True) # image transforms def test_transform(): def _convert_image_to_rgb(im): return im.convert('RGB') return transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), _convert_image_to_rgb, transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) preprocess = test_transform() @torch.no_grad() def denormalize(x, mean, std): # 3, H, W t = x.clone() t.mul_(std).add_(mean) return torch.clamp(t, 0, 1) # Gradio UI def inference(query, *support_text_box_and_files): ''' query: PIL image class_names: list of class names ''' labels = support_text_box_and_files[0::2] support_images = support_text_box_and_files[1::2] print(f"Support images: {support_images}") #first, open the images support_images = [[Image.open(img) for img in imgs] for imgs in support_images if imgs != None] supp_x = [] supp_y = [] for i, support_imgs in enumerate(support_images): #for i, (class_name, support_imgs) in enumerate(zip(class_names, support_images)): if len(support_imgs) == 0: continue for img in support_imgs: x_im = preprocess(img) supp_x.append(x_im) supp_y.append(i) supp_x = torch.stack(supp_x, dim=0).unsqueeze(0).to(device) # (1, n_supp*n_labels, 3, H, W) supp_y = torch.tensor(supp_y).long().unsqueeze(0).to(device) # (1, n_supp*n_labels) query = preprocess(query).unsqueeze(0).unsqueeze(0).to(device) # (1, 3, H, W) print(f"Shape of supp_x: {supp_x.shape}") print(f"Shape of supp_y: {supp_y.shape}") print(f"Shape of query: {query.shape}") with torch.cuda.amp.autocast(True): start_time = time.time() output = model(supp_x, supp_y, query) # (1, 1, n_labels) exec_time = time.time() - start_time probs = output.softmax(dim=-1).detach().cpu().numpy() return {k: float(v) for k, v in zip(labels, probs[0, 0])}, exec_time # DEBUG ##query = Image.open('../labrador-puppy.jpg') #query = Image.open('/Users/hushell/Documents/Dan_tr.png') ##labels = 'dog, cat' #labels = 'girl, sussie' #output = inference(query, labels, n_supp=2) #print(output) title = "# P>M>F few-shot learning pipeline" description = "Short description: We take a ViT-small backbone, which is pre-trained with DINO, and meta-trained on Meta-Dataset; for few-shot classification, we use a ProtoNet classifier. The demo can be viewed as zero-shot since the support set is built by searching images from Google. Note that you may need to play with GIS parameters to get good support examples. Besides, GIS is not very stable as search requests may fail for many reasons (e.g., number of requests reaches the limit of the day). This code is heavely inspired from the original HF space here" article = "
" min_classes = 2 max_classes = 10 def variable_outputs(k): k = int(k) inputs = [] for _ in range(k): inputs.append(gr.Textbox(visible=True)) inputs.append(gr.File(visible=True)) for _ in range(max_classes-k): inputs.append(gr.Textbox(visible=False)) inputs.append(gr.File(visible=False)) return inputs with gr.Blocks() as demo: with gr.Row(): gr.Markdown(title) with gr.Row(): gr.Markdown(description) with gr.Row(): gr.Markdown(article) with gr.Row(): with gr.Column(): query = gr.Image(label="Image to classify", type="pil") num_classes_slider = gr.Slider(minimum=min_classes, maximum=10, value=2, label="Number of classes", step=1) #set_number_classes_btn = gr.Button("Set number of classes") textboxes_and_files = [] for i in range(max_classes): is_visible = (i < 2) t = gr.Textbox(label=f"Class {i+1} name", placeholder=f"Enter class {i+1} name", visible=is_visible) textboxes_and_files.append(t) f = gr.File(label=f"Support image for class {i+1}", type="filepath", visible=is_visible, file_count="multiple") textboxes_and_files.append(f) num_classes_slider.change(variable_outputs, inputs=[num_classes_slider], outputs=textboxes_and_files) run_button = gr.Button("Run Inference") with gr.Column(): output = gr.Label(label="Predicted class probabilities") exec_time = gr.Textbox(label="Execution time (s)") # def run_inference(query, *example_inputs): # # print("len(example_inputs) : ") # print(len(example_inputs)) # # class_names = [example_inputs[i].value for i in range(0, len(example_inputs), 2)] # support_images = [example_inputs[i].value for i in range(1, len(example_inputs), 2)] # return inference(query, class_names, support_images) run_button.click( fn=inference, inputs=[query] + textboxes_and_files, outputs=[output, exec_time] ) # this does nothing it seems demo.examples = [ ["./example_images/2007_000033.jpg", "plane", ["./example_images/2007_000738.jpg", "./example_images/2007_000256.jpg"], "cat", ["./example_images/2007_000528.jpg", "./example_images/2007_000549.jpg"]] ] demo.launch()