File size: 2,832 Bytes
04fbff5 |
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 |
import os
import clip
import torch
import torch.nn as nn
import torch.nn.functional as F
import subprocess
from urllib.request import urlretrieve
from vbench.utils import load_video, load_dimension_info, clip_transform
from tqdm import tqdm
def get_aesthetic_model(cache_folder):
"""load the aethetic model"""
path_to_model = cache_folder + "/sa_0_4_vit_l_14_linear.pth"
if not os.path.exists(path_to_model):
os.makedirs(cache_folder, exist_ok=True)
url_model = (
"https://github.com/LAION-AI/aesthetic-predictor/blob/main/sa_0_4_vit_l_14_linear.pth?raw=true"
)
# download aesthetic predictor
if not os.path.isfile(path_to_model):
try:
print(f'trying urlretrieve to download {url_model} to {path_to_model}')
urlretrieve(url_model, path_to_model) # unable to download https://github.com/LAION-AI/aesthetic-predictor/blob/main/sa_0_4_vit_l_14_linear.pth?raw=true to pretrained/aesthetic_model/emb_reader/sa_0_4_vit_l_14_linear.pth
except:
print(f'unable to download {url_model} to {path_to_model} using urlretrieve, trying wget')
wget_command = ['wget', url_model, '-P', os.path.dirname(path_to_model)]
subprocess.run(wget_command)
m = nn.Linear(768, 1)
s = torch.load(path_to_model)
m.load_state_dict(s)
m.eval()
return m
def laion_aesthetic(aesthetic_model, clip_model, video_list, device):
aesthetic_model.eval()
clip_model.eval()
aesthetic_avg = 0.0
num = 0
video_results = []
for video_path in tqdm(video_list):
images = load_video(video_path)
image_transform = clip_transform(224)
images = image_transform(images)
images = images.to(device)
image_feats = clip_model.encode_image(images).to(torch.float32)
image_feats = F.normalize(image_feats, dim=-1, p=2)
aesthetic_scores = aesthetic_model(image_feats).squeeze()
normalized_aesthetic_scores = aesthetic_scores/10
cur_avg = torch.mean(normalized_aesthetic_scores, dim=0, keepdim=True)
aesthetic_avg += cur_avg.item()
num += 1
video_results.append({'video_path': video_path, 'video_results': cur_avg.item()})
aesthetic_avg /= num
return aesthetic_avg, video_results
def compute_aesthetic_quality(json_dir, device, submodules_list):
vit_path = submodules_list[0]
aes_path = submodules_list[1]
aesthetic_model = get_aesthetic_model(aes_path).to(device)
clip_model, preprocess = clip.load(vit_path, device=device)
video_list, _ = load_dimension_info(json_dir, dimension='aesthetic_quality', lang='en')
all_results, video_results = laion_aesthetic(aesthetic_model, clip_model, video_list, device)
return all_results, video_results
|