|
import torch |
|
import numpy as np |
|
import pandas as pd |
|
import gradio as gr |
|
from io import BytesIO |
|
from PIL import Image as PILIMAGE |
|
|
|
|
|
from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer |
|
import os |
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
model = CLIPModel.from_pretrained("vesteinn/clip-nabirds").to(device) |
|
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") |
|
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32") |
|
|
|
|
|
def load_class_names(dataset_path=''): |
|
names = {} |
|
with open(os.path.join(dataset_path, 'classes.txt')) as f: |
|
for line in f: |
|
pieces = line.strip().split() |
|
class_id = pieces[0] |
|
names[class_id] = ' '.join(pieces[1:]) |
|
|
|
return names |
|
|
|
|
|
def get_labels(): |
|
labels = [] |
|
class_names = load_class_names(".") |
|
for _, name in class_names.items(): |
|
labels.append(f"This is a photo of {name}.") |
|
return labels |
|
|
|
|
|
def encode_text(text): |
|
with torch.no_grad(): |
|
inputs = tokenizer([text], padding=True, return_tensors="pt") |
|
text_encoded = model.get_text_features(**inputs).detach().numpy() |
|
return text_encoded |
|
|
|
|
|
ALL_LABELS = get_labels() |
|
try: |
|
LABEL_FEATURES = np.load("label_features.np") |
|
except: |
|
LABEL_FEATURES = [] |
|
for label in ALL_LABELS: |
|
LABEL_FEATURES.append(encode_text(label)) |
|
LABEL_FEATURES = np.vstack(LABEL_FEATURES) |
|
np.save(open("label_features.np", "wb"), LABEL_FEATURES) |
|
|
|
|
|
def encode_image(image): |
|
image = PILIMAGE.fromarray(image.astype('uint8'), 'RGB') |
|
with torch.no_grad(): |
|
photo_preprocessed = processor(text=None, images=image, return_tensors="pt", padding=True)["pixel_values"] |
|
search_photo_feature = model.get_image_features(photo_preprocessed.to(device)) |
|
search_photo_feature /= search_photo_feature.norm(dim=-1, keepdim=True) |
|
image_encoded = search_photo_feature.cpu().numpy() |
|
return image_encoded |
|
|
|
|
|
def similarity(feature, label_features): |
|
similarities = list((feature @ label_features.T).squeeze(0)) |
|
return similarities |
|
|
|
|
|
def find_best_matches(image): |
|
image_features = encode_image(image) |
|
similarities = similarity(image_features, LABEL_FEATURES) |
|
best_spec = sorted(zip(similarities, range(LABEL_FEATURES.shape[0])), key=lambda x: x[0], reverse=True) |
|
idx = best_spec[0][1] |
|
label = ALL_LABELS[idx] |
|
return label |
|
|
|
examples=[['bj.jpg'],['duckly.jpg'],['some.jpg'],['turdus.jpg'],['seag.jpg'],['thursh.jpg'], ['woodcock.jpeg'],['dipper.jpeg']] |
|
|
|
gr.Interface(fn=find_best_matches, |
|
inputs=[ |
|
gr.inputs.Image(label="Image to classify", optional=False), |
|
], |
|
examples=examples, |
|
theme="grass", |
|
outputs=gr.outputs.Label(), enable_queue=True, title="North American Bird Classifier", |
|
description="This application can classify North American Birds.").launch() |
|
|
|
|
|
|
|
|
|
|
