| 326 | # clip |
| 327 | @torch.no_grad() |
| 328 | def retriev( |
| 329 | model, preprocess, elements: [Image.Image], search_text: str, device |
| 330 | ): |
| 331 | preprocessed_images = [preprocess(image).to(device) for image in elements] |
| 332 | tokenized_text = clip.tokenize([search_text]).to(device) |
| 333 | stacked_images = torch.stack(preprocessed_images) |
| 334 | image_features = model.encode_image(stacked_images) |
| 335 | text_features = model.encode_text(tokenized_text) |
| 336 | image_features /= image_features.norm(dim=-1, keepdim=True) |
| 337 | text_features /= text_features.norm(dim=-1, keepdim=True) |
| 338 | probs = 100.0 * image_features @ text_features.T |
| 339 | return probs[:, 0].softmax(dim=0) |
| 340 | |
| 341 | |
| 342 | def crop_image(annotations, image_like): |