Spaces:
Sleeping
Sleeping
File size: 2,284 Bytes
11e7313 e608f06 11e7313 e608f06 1f43103 5c77158 e608f06 11e7313 e608f06 11e7313 |
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 |
import os,json
from training import clip_dataset
import torch
from torch.utils.data import DataLoader
import open_clip
from datasets import load_dataset
from tqdm import tqdm
import numpy as np
from huggingface_hub import hf_hub_download
import warnings
warnings.filterwarnings("ignore", message=".*QuickGELU mismatch.*")
HF_TOKEN=os.getenv("HF_TOKEN")
MODEL_ID = "PrashantGoyal/findr-clip-ft"
model_path = hf_hub_download(
repo_id=MODEL_ID,
force_download=True,
filename="clip/best.pt",
token=os.getenv("HF_TOKEN")
)
def collate(batch):
img,text=zip(*batch)
return torch.stack(img,0),torch.stack(text,0)
@torch.no_grad
def encode_img(model,processor,tokenizer,split,device):
ds=clip_dataset(split=split,processor=processor,tokenizer=tokenizer)
print('dataset Loaded')
dl=DataLoader(ds,batch_size=4,shuffle=False,num_workers=4,collate_fn=collate)
all_img,all_text=[],[]
for img,text in tqdm(dl,desc=f"Encode {split}"):
img=img.to(device)
text=text.to(device)
img_f=model.encode_image(img)
text_f=model.encode_text(text)
img_f=img_f/img_f.norm(keepdim=True,dim=-1)
text_f=text_f/text_f.norm(keepdim=True,dim=-1)
all_img.append(img_f.cpu())
all_text.append(text_f.cpu())
return torch.cat(all_img),torch.cat(all_text)
def gold_k(sims,k):
ranks = (-sims).argsort(axis=1)
hits = (ranks[:, :k] == np.arange(sims.shape[0])[:,None]).any(axis=1)
return hits.mean()
def main(path='./model/clip/best.pt',arch='ViT-B-32', pretrained='openai'):
device="cuda" if torch.cuda.is_available() else "cpu"
torch.cuda.empty_cache()
model, _, preprocess =open_clip.create_model_and_transforms(arch,pretrained=pretrained,device=device,quick_gelu=True )
tokenizer=open_clip.get_tokenizer(arch)
state=torch.load(model_path,map_location='cuda')['model']
model.load_state_dict(state, strict=False)
model.eval()
print('model loaded')
img_f,text_f=encode_img(model,processor=preprocess,tokenizer=tokenizer,split='test',device=device)
sim=(img_f@text_f.T).numpy()
g1=gold_k(sim,1)
g5=gold_k(sim,5)
g10=gold_k(sim,10)
print(f"Image->Text R@1={g1:.3f} R@5={g5:.3f} R@10={g10:.3f}")
if __name__=="__main__":
main()
|