文件路徑daclip-uir-main/universal-image-restoration/config/daclip-sde/test.py
代碼有部分修改
導(dǎo)包
import argparse
import logging
import os.path
import sys
import time
from collections import OrderedDict
import torchvision.utils as tvutils
import numpy as np
import torch
from IPython import embed
import lpips
import options as option
from models import create_model
sys.path.insert(0, "../../")
import open_clip
import utils as util
from data import create_dataloader, create_dataset
from data.util import bgr2ycbcr
注意open_clip使用的是項(xiàng)目里的代碼,而非環(huán)境里裝的那個(gè)。data、util、option同樣是項(xiàng)目里有的包
聲明
#### options
parser = argparse.ArgumentParser()
parser.add_argument("-opt", type=str, default='options/test.yml', help="Path to options YMAL file.")
opt = option.parse(parser.parse_args().opt, is_train=False)
opt = option.dict_to_nonedict(opt)
配置文件?
設(shè)置配置文件相對(duì)地址options/test.yml
在該配置文件中配置GT和LQ圖像文件地址
datasets:
test1:
name: Test
mode: LQGT
dataroot_GT: C:\Users\86136\Desktop\LQ_test\shadow\GT
dataroot_LQ: C:\Users\86136\Desktop\LQ_test\shadow\LQ
設(shè)置results_root結(jié)果地址,每次計(jì)算結(jié)束這個(gè)地址保存要求記錄的計(jì)算結(jié)果
該目錄下Test文件夾將保存一張GT一張LQ一張復(fù)原圖像??。
不設(shè)置也會(huì)默認(rèn)在項(xiàng)目?jī)?nèi)?daclip-uir-main\results\daclip-sde\universal-ir
#### path
path:
pretrain_model_G: E:\daclip\pretrained\universal-ir.pth
daclip: E:\daclip\pretrained\daclip_ViT-B-32.pt
results_root: C:\Users\86136\Desktop\daclip-uir-main\results\daclip-sde\universal-ir
log:
?
#### mkdir and logger
util.mkdirs(
(
path
for key, path in opt["path"].items()
if not key == "experiments_root"
and "pretrain_model" not in key
and "resume" not in key
)
)
# os.system("rm ./result")
# os.symlink(os.path.join(opt["path"]["results_root"], ".."), "./result")
?報(bào)錯(cuò)執(zhí)行代碼沒(méi)有刪除再創(chuàng)建權(quán)限?我把相關(guān)os操作注釋了,全部保存到result對(duì)我影響不大
加載創(chuàng)建數(shù)據(jù)對(duì)
#### Create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt["datasets"].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set, dataset_opt)
logger.info(
"Number of test images in [{:s}]: {:d}".format(
dataset_opt["name"], len(test_set)
)
)
test_loaders.append(test_loader)
?自定義包含復(fù)原IR-SDE模型的外層類(lèi)model,參考app.py
# load pretrained model by default
model = create_model(opt)
device = model.device
?加載DA-CLIP、IR-SDE
# clip_model, _preprocess = clip.load("ViT-B/32", device=device)
if opt['path']['daclip'] is not None:
clip_model, preprocess = open_clip.create_model_from_pretrained('daclip_ViT-B-32', pretrained=opt['path']['daclip'])
else:
clip_model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
tokenizer = open_clip.get_tokenizer('ViT-B-32')
clip_model = clip_model.to(device)
else是直接使用CLIP的ViT-B-32模型進(jìn)行測(cè)試的代碼。與我測(cè)DA-CLIP無(wú)關(guān)。
想使用的話?目測(cè)要預(yù)先下載對(duì)應(yīng)模型權(quán)重并手動(dòng)修改pretrained為文件地址,否則報(bào)錯(cuò)hf無(wú)法連接
sde = util.IRSDE(max_sigma=opt["sde"]["max_sigma"], T=opt["sde"]["T"], schedule=opt["sde"]["schedule"], eps=opt["sde"]["eps"], device=device)
sde.set_model(model.model)
lpips_fn = lpips.LPIPS(net='alex').to(device)
scale = opt['degradation']['scale']
加載IR-SDE、LPIPS
如果不指定crop_border后續(xù)crop_border=scale
處理并計(jì)算
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt["name"] # path opt['']
logger.info("\nTesting [{:s}]...".format(test_set_name))
test_start_time = time.time()
dataset_dir = os.path.join(opt["path"]["results_root"], test_set_name)
util.mkdir(dataset_dir)
test_results = OrderedDict()
test_results["psnr"] = []
test_results["ssim"] = []
test_results["psnr_y"] = []
test_results["ssim_y"] = []
test_results["lpips"] = []
test_times = []
for i, test_data in enumerate(test_loader):
single_img_psnr = []
single_img_ssim = []
single_img_psnr_y = []
single_img_ssim_y = []
need_GT = False if test_loader.dataset.opt["dataroot_GT"] is None else True
img_path = test_data["GT_path"][0] if need_GT else test_data["LQ_path"][0]
img_name = os.path.splitext(os.path.basename(img_path))[0]
#### input dataset_LQ
LQ, GT = test_data["LQ"], test_data["GT"]
img4clip = test_data["LQ_clip"].to(device)
with torch.no_grad(), torch.cuda.amp.autocast():
image_context, degra_context = clip_model.encode_image(img4clip, control=True)
image_context = image_context.float()
degra_context = degra_context.float()
noisy_state = sde.noise_state(LQ)
model.feed_data(noisy_state, LQ, GT, text_context=degra_context, image_context=image_context)
tic = time.time()
model.test(sde, save_states=False)
toc = time.time()
test_times.append(toc - tic)
visuals = model.get_current_visuals()
SR_img = visuals["Output"]
output = util.tensor2img(SR_img.squeeze()) # uint8
LQ_ = util.tensor2img(visuals["Input"].squeeze()) # uint8
GT_ = util.tensor2img(visuals["GT"].squeeze()) # uint8
suffix = opt["suffix"]
if suffix:
save_img_path = os.path.join(dataset_dir, img_name + suffix + ".png")
else:
save_img_path = os.path.join(dataset_dir, img_name + ".png")
util.save_img(output, save_img_path)
# remove it if you only want to save output images
LQ_img_path = os.path.join(dataset_dir, img_name + "_LQ.png")
GT_img_path = os.path.join(dataset_dir, img_name + "_HQ.png")
util.save_img(LQ_, LQ_img_path)
util.save_img(GT_, GT_img_path)
if need_GT:
gt_img = GT_ / 255.0
sr_img = output / 255.0
crop_border = opt["crop_border"] if opt["crop_border"] else scale
if crop_border == 0:
cropped_sr_img = sr_img
cropped_gt_img = gt_img
else:
cropped_sr_img = sr_img[
crop_border:-crop_border, crop_border:-crop_border
]
cropped_gt_img = gt_img[
crop_border:-crop_border, crop_border:-crop_border
]
psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255)
ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255)
lp_score = lpips_fn(
GT.to(device) * 2 - 1, SR_img.to(device) * 2 - 1).squeeze().item()
test_results["psnr"].append(psnr)
test_results["ssim"].append(ssim)
test_results["lpips"].append(lp_score)
if len(gt_img.shape) == 3:
if gt_img.shape[2] == 3: # RGB image
sr_img_y = bgr2ycbcr(sr_img, only_y=True)
gt_img_y = bgr2ycbcr(gt_img, only_y=True)
if crop_border == 0:
cropped_sr_img_y = sr_img_y
cropped_gt_img_y = gt_img_y
else:
cropped_sr_img_y = sr_img_y[
crop_border:-crop_border, crop_border:-crop_border
]
cropped_gt_img_y = gt_img_y[
crop_border:-crop_border, crop_border:-crop_border
]
psnr_y = util.calculate_psnr(
cropped_sr_img_y * 255, cropped_gt_img_y * 255
)
ssim_y = util.calculate_ssim(
cropped_sr_img_y * 255, cropped_gt_img_y * 255
)
test_results["psnr_y"].append(psnr_y)
test_results["ssim_y"].append(ssim_y)
logger.info(
"img{:3d}:{:15s} - PSNR: {:.6f} dB; SSIM: {:.6f}; LPIPS: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.".format(
i, img_name, psnr, ssim, lp_score, psnr_y, ssim_y
)
)
else:
logger.info(
"img:{:15s} - PSNR: {:.6f} dB; SSIM: {:.6f}.".format(
img_name, psnr, ssim
)
)
test_results["psnr_y"].append(psnr)
test_results["ssim_y"].append(ssim)
else:
logger.info(img_name)
ave_lpips = sum(test_results["lpips"]) / len(test_results["lpips"])
ave_psnr = sum(test_results["psnr"]) / len(test_results["psnr"])
ave_ssim = sum(test_results["ssim"]) / len(test_results["ssim"])
logger.info(
"----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n".format(
test_set_name, ave_psnr, ave_ssim
)
)
if test_results["psnr_y"] and test_results["ssim_y"]:
ave_psnr_y = sum(test_results["psnr_y"]) / len(test_results["psnr_y"])
ave_ssim_y = sum(test_results["ssim_y"]) / len(test_results["ssim_y"])
logger.info(
"----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n".format(
ave_psnr_y, ave_ssim_y
)
)
logger.info(
"----average LPIPS\t: {:.6f}\n".format(ave_lpips)
)
print(f"average test time: {np.mean(test_times):.4f}")
開(kāi)頭往log記錄了相應(yīng)配置文件內(nèi)容,不需要可以注釋。
遍歷測(cè)試數(shù)據(jù)集(test_loaders
)計(jì)算各種評(píng)價(jià)指標(biāo),如峰值信噪比(PSNR)、結(jié)構(gòu)相似性(SSIM)和感知損失(LPIPS)。
在處理過(guò)程中,代碼首先會(huì)創(chuàng)建一個(gè)目錄來(lái)保存測(cè)試結(jié)果。
然后,對(duì)于每個(gè)測(cè)試圖像,代碼會(huì)加載對(duì)應(yīng)的圖像(如果可用),并使用一個(gè)名為clip_model
的模型對(duì)圖像進(jìn)行編碼。
接下來(lái),代碼會(huì)使用一個(gè)名為sde
的隨機(jī)微分方程模型和名為model
的深度學(xué)習(xí)模型來(lái)處理帶有噪聲的圖像,并生成復(fù)原圖像(SR_img)。額可能作者拿了以前做超分的代碼沒(méi)改變量名
在這個(gè)過(guò)程中,text_context
和image_context
被用作模型的輸入,
圖像都會(huì)被保存到之前創(chuàng)建的目錄中。
此外,代碼還會(huì)計(jì)算并記錄每個(gè)圖像的PSNR、SSIM和LPIPS分?jǐn)?shù),并在最后打印出這些分?jǐn)?shù)的平均值。 代碼中還包含了一些用于圖像處理的實(shí)用函數(shù),如util.tensor2img
用于將張量轉(zhuǎn)換為圖像,util.save_img
用于保存圖像,以及util.calculate_psnr
和util.calculate_ssim
用于計(jì)算PSNR和SSIM分?jǐn)?shù)。psnr_y和ssim_y?不用可以把相關(guān)代碼注釋。
最后,代碼還計(jì)算了平均測(cè)試時(shí)間,并將其打印出來(lái)。
結(jié)果
log處理的單張圖像報(bào)錯(cuò)的信息 0是該處理的圖像排序序號(hào),即正在處理第0張圖
24-04-03 17:28:24.697 - INFO: img ?0:_MG_2374_no_shadow - PSNR: 27.779773 dB; SSIM: 0.863140; LPIPS: 0.078669; PSNR_Y: 29.135256 dB; SSIM_Y: 0.869278.
?文章來(lái)源:http://www.zghlxwxcb.cn/news/detail-844831.html
可以給復(fù)原結(jié)果圖加個(gè)后綴方便區(qū)分。文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-844831.html
到了這里,關(guān)于【DA-CLIP】test.py解讀,調(diào)用DA-CLIP和IRSDE模型復(fù)原計(jì)算復(fù)原圖與GT圖SSIM、PSNR、LPIPS的文章就介紹完了。如果您還想了解更多內(nèi)容,請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章,希望大家以后多多支持TOY模板網(wǎng)!