| from typing import List |
|
|
| import gradio as gr |
| import numpy as np |
| import torch |
| from transformers import CLIPProcessor, CLIPModel |
|
|
| IMAGENET_CLASSES_FILE = "imagenet-classes.txt" |
| EXAMPLES = ["dog.jpeg", "car.png"] |
|
|
| MARKDOWN = """ |
| # Zero-Shot Image Classification with MetaCLIP |
| |
| This is the demo for a zero-shot image classification model based on |
| [MetaCLIP](https://github.com/facebookresearch/MetaCLIP), described in the paper |
| [Demystifying CLIP Data](https://arxiv.org/abs/2309.16671) that formalizes CLIP data |
| curation as a simple algorithm. |
| """ |
|
|
|
|
| def load_text_lines(file_path: str) -> List[str]: |
| with open(file_path, 'r') as file: |
| lines = file.readlines() |
| return [line.rstrip() for line in lines] |
|
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model = CLIPModel.from_pretrained("facebook/metaclip-b32-400m").to(device) |
| processor = CLIPProcessor.from_pretrained("facebook/metaclip-b32-400m") |
| imagenet_classes = load_text_lines(IMAGENET_CLASSES_FILE) |
|
|
|
|
| def classify_image(input_image) -> str: |
| inputs = processor( |
| text=imagenet_classes, |
| images=input_image, |
| return_tensors="pt", |
| padding=True).to(device) |
| outputs = model(**inputs) |
| probs = outputs.logits_per_image.softmax(dim=1) |
| class_index = np.argmax(probs.detach().cpu().numpy()) |
| return imagenet_classes[class_index] |
|
|
|
|
| with gr.Blocks() as demo: |
| gr.Markdown(MARKDOWN) |
| with gr.Row(): |
| image = gr.Image(image_mode='RGB', type='pil') |
| output_text = gr.Textbox(label="Output") |
| submit_button = gr.Button("Submit") |
|
|
| submit_button.click(classify_image, inputs=[image], outputs=output_text) |
|
|
| gr.Examples( |
| examples=EXAMPLES, |
| fn=classify_image, |
| inputs=[image], |
| outputs=[output_text], |
| cache_examples=True, |
| run_on_click=True |
| ) |
|
|
| demo.launch(debug=False) |
|
|
