| | import sys, os |
| | import json |
| | from PIL import Image |
| | from tqdm import tqdm |
| | from os.path import dirname, join |
| |
|
| | sys.path.append(dirname(dirname(__file__))) |
| |
|
| | import torch |
| | from transformers import AutoImageProcessor, AutoModel |
| | from transformers import CLIPProcessor, CLIPModel |
| | from transformers import pipeline |
| |
|
| | from data.data import osv5m |
| | from json_stream import streamable_list |
| |
|
| | DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| |
|
| |
|
| | def load_model_clip(): |
| | model = CLIPModel.from_pretrained("laion/CLIP-ViT-L-14-laion2B-s32B-b82K") |
| | processor = CLIPProcessor.from_pretrained("laion/CLIP-ViT-L-14-laion2B-s32B-b82K") |
| | return processor, model.to(DEVICE) |
| |
|
| |
|
| | def load_model_dino(): |
| | model = AutoModel.from_pretrained("facebook/dinov2-base") |
| | processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base") |
| | return processor, model.to(DEVICE) |
| |
|
| |
|
| | def compute_dino(processor, model, x): |
| | inputs = processor(images=x[0], return_tensors="pt", device=DEVICE).to(DEVICE) |
| | outputs = model(**inputs) |
| | last_hidden_states = outputs.last_hidden_state.cpu().numpy() |
| | for i in range(len(x[0])): |
| | yield [last_hidden_states[i].tolist(), x[1][i], x[2][i], x[3][i]] |
| |
|
| |
|
| | def compute_clip(processor, model, x): |
| | inputs = processor(images=x[0], return_tensors="pt", device=DEVICE).to(DEVICE) |
| | features = model.get_image_features(**inputs) |
| | features /= features.norm(dim=-1, keepdim=True) |
| | features = features.cpu().numpy() |
| | for i in range(len(x[0])): |
| | yield [features[i].tolist(), x[1][i], x[2][i], x[3][i]] |
| |
|
| |
|
| | def get_batch(dataset, batch_size): |
| | data, lats, lons, ids = [], [], [], [] |
| | for i in range(len(dataset)): |
| | id, lat, lon = dataset.df.iloc[i] |
| | data.append(Image.open(join(dataset.image_folder, f"{int(id)}.jpg"))) |
| | lats.append(lat) |
| | lons.append(lon) |
| | ids.append(id) |
| | if len(data) == batch_size: |
| | yield data, lats, lons, ids |
| | data, lats, lons, ids = [], [], [], [] |
| |
|
| | if len(data) > 0: |
| | yield data, lats, lons, ids |
| | data, lats, lons, ids = [], [], [], [] |
| |
|
| |
|
| | if __name__ == "__main__": |
| | import argparse |
| |
|
| | parser = argparse.ArgumentParser() |
| | parser.add_argument("--batch_size", type=int, default=256) |
| | parser.add_argument("--compute_features", action="store_true") |
| | parser.add_argument("--compute_nearest", action="store_true") |
| | parser.add_argument("--json_path", default="features") |
| | parser.add_argument("--which", type=str, default="clip", choices=["clip", "dino"]) |
| | args = parser.parse_args() |
| | json_path = join(args.json_path, args.which) |
| |
|
| | os.makedirs(json_path, exist_ok=True) |
| | if args.compute_features: |
| | processor, model = ( |
| | load_model_clip() if args.which == "clip" else load_model_dino() |
| | ) |
| | compute_fn = compute_clip if args.which == "clip" else compute_dino |
| |
|
| | for split in ["test"]: |
| | |
| | json_path_ = join(json_path, f"{split}.json") |
| |
|
| | dataset = OSV5M( |
| | "datasets/osv5m", transforms=None, split=split, dont_split=True |
| | ) |
| |
|
| | @torch.no_grad() |
| | def compute(batch_size): |
| | for data in tqdm( |
| | get_batch(dataset, batch_size), |
| | total=len(dataset) // batch_size, |
| | desc=f"Computing {split} on {args.which}", |
| | ): |
| | features = compute_fn(processor, model, data) |
| | for feature, lat, lon, id in features: |
| | yield feature, lat, lon, id |
| |
|
| | data = streamable_list(compute(args.batch_size)) |
| | json.dump(data, open(json_path_, "w"), indent=4) |
| |
|
| | if args.compute_nearest: |
| | from sklearn.metrics.pairwise import cosine_similarity |
| | import numpy as np |
| |
|
| | train, test = [ |
| | json.load(open(join(json_path, f"{split}.json"), "r")) |
| | for split in ["train", "test"] |
| | ] |
| |
|
| | def get_neighbors(k=10): |
| | for i, test_data in enumerate(tqdm(test)): |
| | feature, lat, lon, id = test_data |
| | features_train = np.stack( |
| | [np.array(train_data[0]) for train_data in train] |
| | ) |
| | cs = np.squeeze( |
| | cosine_similarity(np.expand_dims(feature, axis=0), features_train), |
| | axis=0, |
| | ) |
| | i = np.argsort(cs)[-k:][::-1].tolist() |
| | yield [ |
| | {n: x} |
| | for idx in i |
| | for n, x in zip( |
| | ["feature", "lat", "lon", "id", "distance"], |
| | train[idx] |
| | + [ |
| | cs[idx], |
| | ], |
| | ) |
| | ] |
| |
|
| | data = streamable_list(get_neighbors()) |
| | json.dump(data, open(join(json_path, "nearest.json"), "w"), indent=4) |
| |
|
