分布的差异:KL散度,cross-entropy交叉熵
import os
import torch
import cv2
import argparse
import warnings
import numpy as np
from utils import PSNR, validation, LossNetwork
from model.IAT_main import IAT
from torchvision.transforms import Normalize
import matplotlib.pyplot as plt
parser = argparse.ArgumentParser()
parser.add_argument('--file_name', type=str, default='/home/s3090/zzh/Illumination-Adaptive-Transformer-main/IAT_enhance/Your_Path/')
parser.add_argument('--normalize', type=bool, default=False)
parser.add_argument('--task', type=str, default='enhance', help='Choose from exposure or enhance')
config = parser.parse_args()
# Weights path
exposure_pretrain = r'best_Epoch_exposure.pth'
enhance_pretrain = r'/home/s3090/zzh/Illumination-Adaptive-Transformer-main/IAT_enhance/workdirs/grad3/best_Epoch.pth'
normalize_process = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
## Load Pre-train Weights
model = IAT().cuda()
if config.task == 'exposure':
model.load_state_dict(torch.load(exposure_pretrain))
elif config.task == 'enhance':
model.load_state_dict(torch.load(enhance_pretrain))
else:
warnings.warn('Only could be exposure or enhance')
model.eval()
# 设置视频路径和输出路径
video_path = 'video.mp4'
output_path = 'output_video.mp4'
# 打开视频文件
cap = cv2.VideoCapture(video_path)
# 获取视频的宽度和高度
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# 设置输出视频的编解码器和帧率与输入视频相同
output_codec = cv2.VideoWriter_fourcc(*'mp4v')
output_fps = cap.get(cv2.CAP_PROP_FPS)
# 创建输出视频的写入器
out = cv2.VideoWriter(output_path, output_codec, output_fps, (frame_width, frame_height))
while True:
# 读取视频帧
ret, frame = cap.read()
# 判断是否读取到视频帧
if not ret:
break
# 将BGR图像转换为RGB图像
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# 将NumPy数组转换为PyTorch Tensor
input = torch.from_numpy(frame_rgb / 255.0).float().cuda()
input = input.permute(2, 0, 1).unsqueeze(0)
if config.normalize:
input = normalize_process(input)
# 使用模型进行处理
_, _, enhanced_img = model(input)
enhanced_img = enhanced_img.cpu().squeeze(0).permute(1, 2, 0)
enhanced_img = enhanced_img.detach().numpy()
enhanced_img = (enhanced_img * 255.0).astype(np.uint8)
# 将RGB图像转换为BGR图像
enhanced_img_bgr = cv2.cvtColor(enhanced_img, cv2.COLOR_RGB2BGR)
# 写入增强后的帧到输出视频
out.write(enhanced_img_bgr)