Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,382 Bytes
2cda712 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
# Importing Libraries
import numpy as np
from PIL import Image
import torch
from torchvision import transforms
import os,sys,warnings
warnings.filterwarnings("ignore")
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from functions.ReIQA.options.train_options import TrainOptions
from functions.ReIQA.networks.build_backbone import build_model
from functions.ReIQA.demo_content_aware_feats import run_inference_content
from functions.ReIQA.demo_quality_aware_feats import run_inference_quality
import defaults
class Compute_ReIQA(torch.nn.Module):
def __init__(self,
model_path:str,
device:str
):
"""
Args:
model_path (str): Path to weights of ReIQA feature extractor.
device (str): Device used while computing ReIQA features.
"""
super().__init__()
# Device
if device is None:
self.device = "cuda" if torch.cuda.is_available() else "cpu"
else:
self.device = device
# Paths
self.content_model_path = os.path.join(model_path, "content_aware_r50.pth")
self.quality_model_path = os.path.join(model_path, "quality_aware_r50.pth")
# Arguments
args = TrainOptions().parse()
args.device = self.device
# THE FOLLOWING SETTINGS are set to avoid EXECUTION ISSUES.
strict = False # Generally strict = True in `load_state_dict`
args.feat_dim = 2048 # Added this line to balance dimensions of weights.
# Content Model
self.content_model, _ = build_model(args)
self.content_model = torch.nn.DataParallel(self.content_model)
checkpoint = torch.load(self.content_model_path, map_location=self.device)
self.content_model.load_state_dict(checkpoint['model'], strict=strict)
self.content_model.to(self.device)
self.content_model.eval()
# Quality Model
self.quality_model, _ = build_model(args)
self.quality_model = torch.nn.DataParallel(self.quality_model)
checkpoint = torch.load(self.quality_model_path, map_location=self.device)
self.quality_model.load_state_dict(checkpoint['model'], strict=strict)
self.quality_model.to(self.device)
self.quality_model.eval()
def forward(self, img1, img2):
# ReIQA does not perform L2-normalization and returns non-normalized features
# Quality Features
feat1 = self.quality_model.module.encoder(img1)
feat1 = torch.flatten(feat1, start_dim=1)
feat2 = self.quality_model.module.encoder(img2)
feat2 = torch.flatten(feat2, start_dim=1)
quality = torch.hstack((feat1, feat2))
# Normalization
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
img1 = normalize(img1)
img2 = normalize(img2)
# Content Features
feat1 = self.content_model.module.encoder(img1)
feat1 = torch.flatten(feat1, start_dim=1)
feat2 = self.content_model.module.encoder(img2)
feat2 = torch.flatten(feat2, start_dim=1)
content = torch.hstack((feat1, feat2))
# Content and Quality Fatures
feat = torch.hstack((content, quality))
return feat
# Calling Main function
if __name__ == '__main__':
F = Compute_ReIQA(model_path=os.path.join(defaults.main_feature_ckpts_dir, "feature_extractor_checkpoints/reiqa_feature_extractors"), device="cuda:0")
O = F.forward(torch.randn(1,3,224,224).cuda(), torch.randn(1,3,112,112).cuda())
print (O.shape)
print (torch.linalg.norm(O[0,:2048]), torch.linalg.norm(O[0,2048:4096]), torch.linalg.norm(O[0,4096:6144]), torch.linalg.norm(O[0,6144:8192])) |