Spaces:
Sleeping
Sleeping
| import torch | |
| import open_clip | |
| from PIL import Image | |
| import requests | |
| import json | |
| import gradio as gr | |
| import pandas as pd | |
| from io import BytesIO | |
| import os | |
| # Load the Amazon taxonomy from a JSON file | |
| with open("amazon.json", "r") as f: | |
| AMAZON_TAXONOMY = json.load(f) | |
| base_model_name = "ViT-B-16" | |
| model_base, _, preprocess_base = open_clip.create_model_and_transforms(base_model_name) | |
| tokenizer_base = open_clip.get_tokenizer(base_model_name) | |
| model_name_B = "hf-hub:Marqo/marqo-ecommerce-embeddings-B" | |
| model_B, _, preprocess_B = open_clip.create_model_and_transforms(model_name_B) | |
| tokenizer_B = open_clip.get_tokenizer(model_name_B) | |
| model_name_L = "hf-hub:Marqo/marqo-ecommerce-embeddings-L" | |
| model_L, _, preprocess_L = open_clip.create_model_and_transforms(model_name_L) | |
| tokenizer_L = open_clip.get_tokenizer(model_name_L) | |
| models = [base_model_name, model_name_B, model_name_L] | |
| taxonomy_cache = {} | |
| for model in models: | |
| with open(f'{model.split("/")[-1]}.json', "r") as f: | |
| taxonomy_cache[model] = json.load(f) | |
| def cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: | |
| numerator = (a * b).sum(dim=-1) | |
| denominator = torch.linalg.norm(a, ord=2, dim=-1) * torch.linalg.norm( | |
| b, ord=2, dim=-1 | |
| ) | |
| return 0.5 * (numerator / denominator + 1.0) | |
| class BeamPath: | |
| def __init__(self, path: list, cumulative_score: float, current_layer: dict | list): | |
| self.path = path | |
| self.cumulative_score = cumulative_score | |
| self.current_layer = current_layer | |
| def __repr__(self): | |
| return f"BeamPath(path={self.path}, cumulative_score={self.cumulative_score})" | |
| def _compute_similarities(classes: list, base_embedding: torch.Tensor, cache_key: str): | |
| text_features = torch.tensor( | |
| [taxonomy_cache[cache_key][class_name] for class_name in classes] | |
| ) | |
| similarities = cosine_similarity(base_embedding, text_features) | |
| return similarities.cpu().numpy() | |
| def map_taxonomy( | |
| base_image: Image.Image, | |
| taxonomy: dict, | |
| model, | |
| tokenizer, | |
| preprocess_val, | |
| cache_key, | |
| beam_width: int = 3, | |
| ) -> tuple[list[tuple[str, float]], float]: | |
| image_tensor = preprocess_val(base_image).unsqueeze(0) | |
| with torch.no_grad(), torch.cuda.amp.autocast(): | |
| base_embedding = model.encode_image(image_tensor, normalize=True) | |
| initial_path = BeamPath(path=[], cumulative_score=0.0, current_layer=taxonomy) | |
| beam = [initial_path] | |
| final_paths = [] | |
| is_first = True | |
| while beam: | |
| candidates = [] | |
| candidate_entries = [] | |
| for beam_path in beam: | |
| layer = beam_path.current_layer | |
| if isinstance(layer, dict): | |
| classes = list(layer.keys()) | |
| elif isinstance(layer, list): | |
| classes = layer | |
| if classes == []: | |
| final_paths.append(beam_path) | |
| continue | |
| else: | |
| final_paths.append(beam_path) | |
| continue | |
| # current_path_class_names = [class_name for class_name, _ in beam_path.path] | |
| for class_name in classes: | |
| candidate_string = class_name | |
| if isinstance(layer, dict): | |
| next_layer = layer[class_name] | |
| else: | |
| next_layer = None | |
| candidate_entries.append( | |
| (candidate_string, class_name, beam_path, next_layer) | |
| ) | |
| if not candidate_entries: | |
| break | |
| candidate_strings = [ | |
| candidate_string for candidate_string, _, _, _ in candidate_entries | |
| ] | |
| similarities = _compute_similarities( | |
| candidate_strings, base_embedding, cache_key | |
| ) | |
| for (candidate_string, class_name, beam_path, next_layer), similarity in zip( | |
| candidate_entries, similarities | |
| ): | |
| new_path = beam_path.path + [(class_name, float(similarity))] | |
| new_cumulative_score = beam_path.cumulative_score + similarity | |
| candidate = BeamPath( | |
| path=new_path, | |
| cumulative_score=new_cumulative_score, | |
| current_layer=next_layer, | |
| ) | |
| candidates.append(candidate) | |
| from collections import defaultdict | |
| by_parents = defaultdict(list) | |
| for candidate in candidates: | |
| by_parents[candidate.path[0][0]].append(candidate) | |
| beam = [] | |
| for parent in by_parents: | |
| children = by_parents[parent] | |
| children.sort( | |
| key=lambda x: x.cumulative_score / len(x.path) + x.path[-1][1], | |
| reverse=True, | |
| ) | |
| if is_first: | |
| beam.extend(children) | |
| else: | |
| beam.extend(children[:beam_width]) | |
| is_first = False | |
| all_paths = beam + final_paths | |
| if all_paths: | |
| all_paths.sort(key=lambda x: x.cumulative_score / len(x.path), reverse=True) | |
| best_path = all_paths[0] | |
| return best_path.path, float(best_path.cumulative_score) | |
| else: | |
| return [], 0.0 | |
| # Function to classify image and map taxonomy | |
| def classify_image( | |
| image_input: Image.Image | None, | |
| image_url: str | None, | |
| model_size: str, | |
| beam_width: int, | |
| ): | |
| if image_input is not None: | |
| image = image_input | |
| elif image_url: | |
| # Try to get image from URL | |
| try: | |
| response = requests.get(image_url) | |
| image = Image.open(BytesIO(response.content)).convert("RGB") | |
| except Exception as e: | |
| return pd.DataFrame({"Error": [str(e)]}) | |
| else: | |
| return pd.DataFrame( | |
| { | |
| "Error": [ | |
| "Please provide an image, an image URL, or select an example image" | |
| ] | |
| } | |
| ) | |
| # Select the model, tokenizer, and preprocess | |
| if model_size == "marqo-ecommerce-embeddings-L": | |
| key = "hf-hub:Marqo/marqo-ecommerce-embeddings-L" | |
| model = model_L | |
| preprocess_val = preprocess_L | |
| tokenizer = tokenizer_L | |
| elif model_size == "marqo-ecommerce-embeddings-B": | |
| key = "hf-hub:Marqo/marqo-ecommerce-embeddings-B" | |
| model = model_B | |
| preprocess_val = preprocess_B | |
| tokenizer = tokenizer_B | |
| elif model_size == "openai-ViT-B-16": | |
| key = "ViT-B-16" | |
| model = model_base | |
| preprocess_val = preprocess_base | |
| tokenizer = tokenizer_base | |
| else: | |
| return pd.DataFrame({"Error": ["Invalid model size"]}) | |
| path, cumulative_score = map_taxonomy( | |
| base_image=image, | |
| taxonomy=AMAZON_TAXONOMY, | |
| model=model, | |
| tokenizer=tokenizer, | |
| preprocess_val=preprocess_val, | |
| cache_key=key, | |
| beam_width=beam_width, | |
| ) | |
| output = [] | |
| for idx, (category, score) in enumerate(path): | |
| level = idx + 1 | |
| output.append({"Level": level, "Category": category, "Score": score}) | |
| df = pd.DataFrame(output) | |
| return df | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Image Classification with Taxonomy Mapping") | |
| gr.Markdown( | |
| "## How to use this app\n\nThis app compares [Marqo's Ecommerce embeddings](https://huggingface.co/collections/Marqo/marqo-ecommerce-embeddings-66f611b9bb9d035a8d164fbb) to OpenAI's ViT-B-16 CLIP model for Ecommerce taxonomy mapping. A beam search is used to find the correct classification in the taxonomy. The original OpenAI CLIP models perform very poorly on Ecommerce data." | |
| ) | |
| gr.Markdown( | |
| "Upload an image, provide an image URL, or select an example image, select the model size, and get the taxonomy mapping. The taxonomy is based on the Amazon product taxonomy." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_input = gr.Image(type="pil", label="Upload Image", height=300) | |
| image_url_input = gr.Textbox( | |
| lines=1, placeholder="Image URL", label="Image URL" | |
| ) | |
| gr.Markdown("### Or select an example image:") | |
| # Get example images from 'images' folder | |
| example_images_folder = "images" | |
| example_image_paths = [ | |
| os.path.join(example_images_folder, img) | |
| for img in os.listdir(example_images_folder) | |
| ] | |
| gr.Examples( | |
| examples=[[img_path] for img_path in example_image_paths], | |
| inputs=image_input, | |
| label="Example Images", | |
| examples_per_page=100, | |
| ) | |
| with gr.Column(): | |
| model_size_input = gr.Radio( | |
| choices=[ | |
| "marqo-ecommerce-embeddings-L", | |
| "marqo-ecommerce-embeddings-B", | |
| "openai-ViT-B-16", | |
| ], | |
| label="Model", | |
| value="marqo-ecommerce-embeddings-L", | |
| ) | |
| beam_width_input = gr.Number( | |
| label="Beam Width", value=5, minimum=1, step=1 | |
| ) | |
| classify_button = gr.Button("Classify") | |
| output_table = gr.Dataframe(headers=["Level", "Category", "Score"]) | |
| classify_button.click( | |
| fn=classify_image, | |
| inputs=[image_input, image_url_input, model_size_input, beam_width_input], | |
| outputs=output_table, | |
| ) | |
| demo.launch() | |