Spaces:
Running
Running
| import pickle | |
| import os | |
| from sklearn.neighbors import NearestNeighbors | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image | |
| data_root = 'https://ai-vision-public-datasets.s3.eu.cloud-object-storage.appdomain.cloud/DomainNet' | |
| feat_dir = 'brad_feats' | |
| domains = ['sketch', 'painting', 'clipart', 'real'] | |
| shots = '-1' | |
| num_nn = 20 | |
| search_domain = 'all' | |
| num_results_per_domain = 5 | |
| src_data_dict = {} | |
| class_list = [] | |
| if search_domain == 'all': | |
| for d in domains: | |
| with open(os.path.join(feat_dir, f'dst_{d}_{shots}.pkl'), 'rb') as fp: | |
| src_data = pickle.load(fp) | |
| if class_list == []: | |
| for p in src_data[0]: | |
| cl = p.split('/')[-2] | |
| if cl not in class_list: | |
| class_list.append(cl) | |
| src_nn_fit = NearestNeighbors(n_neighbors=num_results_per_domain, algorithm='auto', n_jobs=-1).fit(src_data[1]) | |
| src_data_dict[d] = (src_data,src_nn_fit) | |
| else: | |
| with open(os.path.join(feat_dir, f'dst_{search_domain}_{shots}.pkl'), 'rb') as fp: | |
| src_data = pickle.load(fp) | |
| src_nn_fit = NearestNeighbors(n_neighbors=num_results_per_domain, algorithm='auto', n_jobs=-1).fit(src_data[1]) | |
| src_data_dict[search_domain] = (src_data,src_nn_fit) | |
| dst_data_dict = {} | |
| min_len = 1e10 | |
| for d in domains: | |
| with open(os.path.join(feat_dir, f'src_{d}_{shots}.pkl'), 'rb') as fp: | |
| dest_data = pickle.load(fp) | |
| dst_data_dict[d] = ({cl: ([],[]) for cl in class_list},dest_data[1]) | |
| for c, p in enumerate(dest_data[0]): | |
| cl = p.split('/')[-2] | |
| dst_data_dict[d][0][cl][0].append(p) | |
| dst_data_dict[d][0][cl][1].append(c) | |
| for cl in class_list: | |
| min_len = min(min_len, len(dst_data_dict[d][0][cl])) | |
| def query(query_index, query_domain, cl): | |
| dst_data = dst_data_dict[query_domain] | |
| dst_img_path = os.path.join(data_root, dst_data[0][cl][0][query_index]) | |
| query_index = dst_data[0][cl][1][query_index] | |
| img_paths = [dst_img_path] | |
| q_cl = dst_img_path.split('/')[-2] | |
| captions = [f'Query: {q_cl}'.title()] | |
| for s_domain, s_data in src_data_dict.items(): | |
| _, top_n_matches_ids = s_data[1].kneighbors(dst_data[1][query_index:query_index+1]) | |
| top_n_labels = s_data[0][2][top_n_matches_ids][0] | |
| src_img_pths = [os.path.join(data_root, s_data[0][0][ix]) for ix in top_n_matches_ids[0]] | |
| img_paths += src_img_pths | |
| for p in src_img_pths: | |
| src_cl = p.split('/')[-2] | |
| src_file = p.split('/')[-1] | |
| captions.append(src_cl.title()) | |
| # print(img_paths) | |
| return tuple([p for p in img_paths])+ tuple(captions) | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.Markdown('# Unsupervised Domain Generalization by Learning a Bridge Across Domains') | |
| gr.Markdown('This demo showcases the cross-domain retrieval capabilities of our self-supervised cross domain training as presented @CVPR 2022. For details please refer to [the paper](https://openaccess.thecvf.com/content/CVPR2022/papers/Harary_Unsupervised_Domain_Generalization_by_Learning_a_Bridge_Across_Domains_CVPR_2022_paper.pdf)') | |
| gr.Markdown('The model is trained in an unsupervised manner on all domains without class labels. The labels are displayed to indicate retrieval success/failure.') | |
| gr.Markdown('## Instructions:') | |
| gr.Markdown('Select a query domain and a class from the drop-down menus and select any random image index from the domain using the slider below, then press the "Run" button. The query image and the retrieved results from each of the four domains, along with the class label will be presented.') | |
| gr.Markdown('## Select Query Domain: ') | |
| gr.Markdown('# Query Image: \t\t\t\t') | |
| # domain_drop = gr.Dropdown(domains) | |
| # cl_drop = gr.Dropdown(class_list) | |
| # domain_select_button = gr.Button("Select Domain") | |
| # slider = gr.Slider(0, min_len) | |
| # slider = gr.Slider(0, 10000) | |
| with gr.Row(): | |
| with gr.Column(): | |
| domain_drop = gr.Dropdown(domains, label='Domain') | |
| cl_drop = gr.Dropdown(class_list, label='Query Class') | |
| slider = gr.Slider(0, 100, label='Query image selector slider') | |
| # gr.Markdown('\t') | |
| # gr.Markdown('\t') | |
| # gr.Markdown('\t') | |
| with gr.Column(): | |
| src_cap = gr.Label() | |
| src_img = gr.Image() | |
| image_button = gr.Button("Run") | |
| out_images = [] | |
| out_captions = [] | |
| for d in domains: | |
| gr.Markdown(f'# Retrieved Images from {d.title()} Domain:') | |
| with gr.Row(): | |
| for _ in range(num_results_per_domain): | |
| with gr.Column(): | |
| out_captions.append(gr.Label()) | |
| out_images.append(gr.Image()) | |
| image_button.click(query, inputs=[slider, domain_drop, cl_drop], outputs=[src_img]+out_images +[src_cap]+ out_captions) | |
| demo.launch(share=True) | |