|
|
import os |
|
|
from tqdm import tqdm |
|
|
import time |
|
|
from utils.config import get_args |
|
|
from pathlib import Path |
|
|
import pickle |
|
|
import torch |
|
|
import open_clip |
|
|
from open_clip import tokenizer |
|
|
import numpy as np |
|
|
|
|
|
def load_clip(): |
|
|
print(f'[INFO] loading CLIP model...') |
|
|
model, _, _ = open_clip.create_model_and_transforms("ViT-H-14", pretrained="laion2b_s32b_b79k") |
|
|
model.cuda() |
|
|
model.eval() |
|
|
print(f'[INFO]', ' finish loading CLIP model...') |
|
|
return model |
|
|
|
|
|
def extract_text_feature(descriptions, clip_model, target_path): |
|
|
text_tokens = tokenizer.tokenize(descriptions).cuda() |
|
|
with torch.no_grad(): |
|
|
text_features = clip_model.encode_text(text_tokens).float() |
|
|
text_features /= text_features.norm(dim=-1, keepdim=True) |
|
|
text_features = text_features.cpu().numpy() |
|
|
|
|
|
text_features_dict = {} |
|
|
for i, description in enumerate(descriptions): |
|
|
text_features_dict[description] = text_features[i] |
|
|
|
|
|
np.save(os.path.join(target_path, "text_features.npy"), text_features_dict) |
|
|
return text_features_dict |
|
|
|
|
|
|
|
|
clip_model = load_clip() |
|
|
|
|
|
def execute_commands(commands_list, command_type, process_num): |
|
|
print('====> Start', command_type) |
|
|
from multiprocessing import Pool |
|
|
pool = Pool(process_num) |
|
|
for _ in tqdm(pool.imap_unordered(os.system, commands_list), total=len(commands_list)): |
|
|
pass |
|
|
pool.close() |
|
|
pool.join() |
|
|
pool.terminate() |
|
|
print('====> Finish', command_type) |
|
|
|
|
|
|
|
|
def parallel_compute(general_command, command_name, resource_type, cuda_list, seq_name_list): |
|
|
cuda_num = len(cuda_list) |
|
|
|
|
|
if resource_type == 'cuda': |
|
|
commands = [] |
|
|
for i, cuda_id in enumerate(cuda_list): |
|
|
process_seq_name = seq_name_list[i::cuda_num] |
|
|
if len(process_seq_name) == 0: |
|
|
continue |
|
|
process_seq_name = '+'.join(process_seq_name) |
|
|
command = f'CUDA_VISIBLE_DEVICES={cuda_id} {general_command % process_seq_name}' |
|
|
commands.append(command) |
|
|
execute_commands(commands, command_name, cuda_num) |
|
|
elif resource_type == 'cpu': |
|
|
commands = [] |
|
|
for seq_name in seq_name_list: |
|
|
commands.append(f'{general_command} --seq_name {seq_name}') |
|
|
execute_commands(commands, command_name, cuda_num) |
|
|
|
|
|
|
|
|
def main(args): |
|
|
CUDA_LIST = args.devices |
|
|
config = "itw" |
|
|
cropformer_path = args.cropformer_path |
|
|
|
|
|
root = f'data/{config}/processed' |
|
|
image_path_pattern = 'color/*.jpg' |
|
|
|
|
|
seq_name_list = os.listdir(root) |
|
|
|
|
|
with open("labels.txt", "r") as f: |
|
|
labels = f.read().split(";") |
|
|
extract_text_feature(labels, clip_model, f"{root}/{seq_name_list[0]}") |
|
|
|
|
|
|
|
|
parallel_compute(f'PYTHONPATH=. python semantics/itw-ovq.py --config {config}', 'get text labels', 'cpu', CUDA_LIST, seq_name_list) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
args = get_args() |
|
|
main(args) |
|
|
|
