|
|
|
|
|
import os |
|
|
import glob |
|
|
import json |
|
|
import tqdm |
|
|
import natsort |
|
|
import random |
|
|
|
|
|
from PIL import Image |
|
|
|
|
|
import numpy as np |
|
|
|
|
|
import torch |
|
|
from torch.utils.data import Dataset, DataLoader |
|
|
|
|
|
import clip |
|
|
|
|
|
from torchvision import models |
|
|
|
|
|
from config import config |
|
|
|
|
|
|
|
|
class loading_img(Dataset): |
|
|
def __init__(self, img_list): |
|
|
self.img_list = img_list |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.img_list) |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
return preprocess(Image.open(self.img_list[idx])) |
|
|
|
|
|
|
|
|
|
|
|
def select_frames(folder_list, preprocess, resnet18_pretrained): |
|
|
for folder in folder_list: |
|
|
img_list = natsort.natsorted(glob.glob(f"{folder}/*.jpg")) |
|
|
img_feats = [] |
|
|
|
|
|
img_set = loading_img(img_list) |
|
|
img_loader = DataLoader(img_set, batch_size=64, shuffle=False, num_workers=16) |
|
|
|
|
|
for imgtensor in img_loader: img_feats.append(imgtensor) |
|
|
img_feats = torch.concat(img_feats, dim=0).to(device) |
|
|
|
|
|
with torch.no_grad(): |
|
|
featuremap = resnet18_pretrained(img_feats) |
|
|
frame_num = featuremap.shape[0] |
|
|
|
|
|
dist_list = [] |
|
|
for img_feat in featuremap: dist_list.append(torch.mean(torch.sqrt((featuremap-img_feat)**2), dim=-1)) |
|
|
dist_list = torch.concat(dist_list).reshape(frame_num, frame_num) |
|
|
|
|
|
idx_list = [_ for _ in range(frame_num)] |
|
|
loop_idx = 0 |
|
|
out_frames = [] |
|
|
|
|
|
output_results = [] |
|
|
while len(idx_list) > 5: |
|
|
dist_idx = idx_list.pop(0) |
|
|
|
|
|
data = dist_list[dist_idx, idx_list].softmax(dim=-1) |
|
|
mu, std = torch.mean(data), torch.std(data) |
|
|
pop_idx_list = torch.where(data < mu-std*(np.exp(1-loop_idx/config.divlam)))[0].detach().cpu().numpy() |
|
|
result = list(np.array(idx_list)[pop_idx_list]) |
|
|
result.append(dist_idx) |
|
|
output_results.append(result) |
|
|
|
|
|
num_picks = 18 |
|
|
if len(result) > num_picks: |
|
|
idx_result_list = sorted(random.sample(result, num_picks)) |
|
|
img_list = np.array(img_list) |
|
|
idx_result_list = np.array(idx_result_list) |
|
|
out_frames.extend(img_list[idx_result_list]) |
|
|
else: |
|
|
idx_result_list = sorted(result) |
|
|
img_list = np.array(img_list) |
|
|
idx_result_list = np.array(idx_result_list) |
|
|
out_frames.extend(img_list[idx_result_list]) |
|
|
|
|
|
loop_idx += 1 |
|
|
|
|
|
for pop_idx in reversed(pop_idx_list): idx_list.pop(pop_idx) |
|
|
|
|
|
return out_frames, output_results |
|
|
|
|
|
|
|
|
|
|
|
random.seed(10) |
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
resnet18_pretrained = models.resnet18(pretrained=True).to(device) |
|
|
resnet18_pretrained.fc = torch.nn.Identity() |
|
|
resnet18_pretrained.avgpool = torch.nn.Identity() |
|
|
resnet18_pretrained.eval() |
|
|
|
|
|
model, preprocess = clip.load("ViT-B/32", device=device) |
|
|
|
|
|
objs_acts = config.f_path |
|
|
questions = config.q_path |
|
|
|
|
|
questions = [json.loads(q) for q in open(os.path.expanduser(questions), "r")] |
|
|
objs_acts = [json.loads(q) for q in open(os.path.expanduser(objs_acts), "r")] |
|
|
|
|
|
answer_path = os.path.expanduser(config.a_path) |
|
|
os.makedirs(os.path.dirname(answer_path), exist_ok=True) |
|
|
ans_file = open(answer_path, "w") |
|
|
|
|
|
output_results = [] |
|
|
for question in tqdm.tqdm(questions): |
|
|
test_token = True |
|
|
|
|
|
for objs_act in objs_acts: |
|
|
if objs_act['q_uid'] == question['q_uid']: |
|
|
question['Object'] = objs_act["Activity"] |
|
|
question['Activity'] = objs_act["Activity"] |
|
|
|
|
|
folder_list = glob.glob(f"{config.img_folder}/{question['q_uid']}/") |
|
|
out_frames, output_result = select_frames(folder_list, preprocess, resnet18_pretrained) |
|
|
output_results.append(output_result) |
|
|
question['filepath'] = out_frames |
|
|
|
|
|
ans_file.write(json.dumps(question) + "\n") |
|
|
test_token = False |
|
|
break |
