Spaces:
Sleeping
Sleeping
File size: 4,158 Bytes
d464b17 bd0098a d464b17 bd0098a d464b17 f572edc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np
import gradio as gr
from facenet_pytorch import InceptionResnetV1, MTCNN
from datasets import load_dataset
from tqdm import tqdm
import wikipedia
dataset = load_dataset("tonyassi/celebrity-1000", split="train")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
transform = transforms.Compose([
transforms.Resize((160, 160)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])
])
class HF_CelebDataset(Dataset):
def __init__(self, hf_dataset, transform=None):
self.dataset = hf_dataset
self.transform = transform
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
img = self.dataset[idx]['image']
label = self.dataset[idx]['label']
if img.mode != 'RGB':
img = img.convert('RGB')
if self.transform:
img = self.transform(img)
return img, label
subset_dataset = HF_CelebDataset(dataset, transform)
loader = DataLoader(subset_dataset, batch_size=64, shuffle=False, num_workers=2)
facenet = InceptionResnetV1(pretrained='vggface2').eval().to(device)
mtcnn = MTCNN(image_size=160, margin=20, device=device)
embeddings, labels = [], []
with torch.no_grad():
for imgs, lbls in tqdm(loader):
imgs = imgs.to(device)
emb = facenet(imgs)
emb = F.normalize(emb, p=2, dim=1)
embeddings.append(emb.cpu())
labels.extend(lbls.numpy())
embeddings = torch.cat(embeddings)
labels = np.array(labels)
unique_labels = np.unique(labels)
avg_embs, celeb_labels, celeb_names = [], [], []
for lbl in unique_labels:
idxs = np.where(labels == lbl)[0]
mean_emb = embeddings[idxs].mean(dim=0)
mean_emb = F.normalize(mean_emb, p=2, dim=0)
avg_embs.append(mean_emb)
celeb_labels.append(lbl)
celeb_names.append(dataset.features['label'].names[lbl])
celeb_embeddings = torch.stack(avg_embs).to(device)
def find_most_similar(user_img, model, celeb_embeddings, celeb_names, top_k=1):
face = mtcnn(user_img)
if face is None:
face = transform(user_img)
img_tensor = face.unsqueeze(0).to(device)
with torch.no_grad():
user_emb = model(img_tensor)
user_emb = F.normalize(user_emb, p=2, dim=1)
similarity = torch.matmul(user_emb, celeb_embeddings.T)
topk_vals, topk_idx = torch.topk(similarity, top_k, dim=1)
top_labels = [celeb_names[i] for i in topk_idx[0]]
top_scores = [float(topk_vals[0][i]) for i in range(top_k)]
return list(zip(top_labels, top_scores))
def get_wikipedia_info(name):
try:
summary = wikipedia.summary(name, sentences=3, auto_suggest=False, redirect=True)
return summary
except wikipedia.exceptions.DisambiguationError as e:
return f"Multiple results found for {name}: {e.options[:3]}"
except wikipedia.exceptions.PageError:
return f"No information found for {name}."
except Exception as e:
return f"Error fetching info: {e}"
def gradio_find(user_img):
top_result = find_most_similar(user_img, facenet, celeb_embeddings, celeb_names, top_k=1)[0]
name, score = top_result
lbl = celeb_labels[celeb_names.index(name)]
idxs = np.where(labels == lbl)[0]
result_imgs = []
for i in idxs[:6]:
img = dataset[int(i)]['image'].resize((128, 128))
caption = f"{name}"
result_imgs.append((img, caption))
wiki_text = get_wikipedia_info(name)
return result_imgs, wiki_text
with gr.Blocks() as demo:
gr.Markdown("## Which Celebrity Do You Look Like?")
with gr.Row():
img_input = gr.Image(type="pil", label="Upload your face")
gallery = gr.Gallery(label="Similar celebrity", columns=3)
info_box = gr.Textbox(label="Wikipedia Info", lines=10)
btn = gr.Button("Find Similar")
btn.click(fn=gradio_find, inputs=img_input, outputs=[gallery, info_box])
gr.close_all()
demo.launch(share=True, inbrowser=True, prevent_thread_lock=True)
if __name__ == "__main__":
demo.launch() |