turns-00085.parquet:54848
71dfcfea5145fbd111425f46degenerate_repetitionAbsentFinal dense release
USER
请参照下述代码中读取npy文件的方式：
# -*- coding: utf-8 -*-
"""
@Author : zhwzhong
@License : (C) Copyright 2013-2018, hit
@Contact : <PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>
@Software: PyCharm
@File : quick_test.py
@Time : 2021/8/31 20:32
@Desc :
"""
import glob
import math
import tqdm
import numpy as np
import torch
import utility
from option import args
from data import get_dataloader
from importlib import import_module
args.scale = 8
args.down_type = 'nearest'
device = torch.device('cpu' if args.cpu else 'cuda')
module = import_module('models.' + args.model_name.lower())
model = module.make_model(args).to(device)
model = torch.nn.parallel.DataParallel(model, device_ids=list(range(args.num_gpus)))
print(utility.get_parameter_number(model))
device_id = torch.cuda.current_device()
load_name = './pre_trained/net_{}_x{}.pth'.format(args.down_type, args.scale)
print(load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage.cuda(device_id))
model.load_state_dict(checkpoint['state'])
model.eval()
test_set = ['Middlebury', 'Lu', 'test', 'Sintel'] # , 'Middlebury', 'test', 'Sintel'
for test_name in test_set:
sum_rmse = []
for gt_name in tqdm.tqdm(sorted(glob.glob('./test_data/{}/gt/*.npy'.format(test_name)))):
gt_img = np.load(gt_name)
rgb_img = np.load(gt_name.replace('gt', 'rgb'))
# Following DKN, we use bicubic in PIL to degrade GT image (for bicubic), and crop the border
# before calculate the RMSE values. (reference: https://github.com/cvlab-yonsei/dkn/issues/1)
module = max(int(math.pow(2, 1 + args.num_pyramid)), args.scale)
tmp_gt = utility.mod_crop(gt_img, modulo=module)
if args.down_type == 'nearest':
tmp_gt = (tmp_gt - np.min(tmp_gt)) / (np.max(tmp_gt) - np.min(tmp_gt))
lr_img = utility.get_lowers(tmp_gt, factor=args.scale, mode=args.down_direction)
else:
tmp_gt = (tmp_gt - np.min(tmp_gt)) / (np.max(tmp_gt) - np.min(tmp_gt))
lr_img = utility.get_lowers(tmp_gt, factor=args.scale, mode='bicubic')
lr_up = utility.get_lowers(lr_img, factor=1 / args.scale, mode='bicubic')
lr_img, gt_img = np.expand_dims(lr_img, 0), np.expand_dims(gt_img, 0)
lr_up = np.expand_dims(lr_up, 0)
if args.guide_channels == 1:
rgb_img = np.expand_dims(utility.rgb2gray(rgb_img), 2)
rgb_img = np.float32(np.transpose(rgb_img, axes=(2, 0, 1))) / 255.
gt_img, rgb_img = utility.mod_crop(gt_img, modulo=module), utility.mod_crop(rgb_img, modulo=module)
lr_img, lr_up, gt_img, rgb_img = utility.np_to_tensor(lr_img, lr_up, gt_img, rgb_img)
lr_img, lr_up, gt_img, rgb_img = lr_img.unsqueeze(0), lr_up.unsqueeze(0), gt_img.unsqueeze(0), rgb_img.unsqueeze(0)
lr_img, lr_up, gt_img, rgb_img = lr_img.to(device), lr_up.to(device), gt_img.to(device), rgb_img.to(device)
out = model(lr=lr_img.contiguous(), rgb=rgb_img.contiguous(), lr_up=lr_up.contiguous())[-1]
if test_name == 'test':
mul_ratio = 100
elif test_name == 'Sintel':
mul_ratio = 255
else:
mul_ratio = 1
rmse, _ = utility.root_mean_sqrt_error(im_pred=out.contiguous(), im_true=gt_img.contiguous(), border=6, mul_ratio=mul_ratio, is_train=False)
sum_rmse.append(rmse)
print('{}: {:.2f}'.format(test_name, np.mean(sum_rmse)))
修改下述代码：
import os
import cv2
import glob
import torch
import numpy as np
from ahmf import AHMF
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
def mod_crop(*imgs, modulo):
def crop(img):
if img.ndim == 2:
h, w = img.shape
crop_h, crop_w = h % modulo, w % modulo
return img[crop_h // 2 : h - (crop_h - crop_h // 2),
crop_w // 2 : w - (crop_w - crop_w // 2)]
elif img.ndim == 3:
c, h, w = img.shape
crop_h, crop_w = h % modulo, w % modulo
return img[:, crop_h // 2 : h - (crop_h - crop_h // 2),
crop_w // 2 : w - (crop_w - crop_w // 2)]
else:
raise ValueError(f"Unsupported image dimension: {img.ndim}")
out = [crop(img) for img in imgs]
return out if len(out) > 1 else out[0]
def np_to_tensor(*imgs, input_range=1.0, output_range=1.0):
def _to_tensor(img):
arr = img.astype(np.float32)
if arr.ndim == 2:
arr = np.expand_dims(arr, 0)
tensor = torch.from_numpy(arr).float()
tensor *= output_range / input_range
return tensor
out = [_to_tensor(img) for img in imgs]
return out if len(out) > 1 else out[0]
def quantize(img, rgb_range):
pixel_range = 255 / rgb_range
return img.mul(pixel_range).clamp(0, 255).round().div(pixel_range)
def imresize(img, scale_factor):
if img.ndim == 2:
h, w = img.shape
new_h, new_w = int(h * scale_factor), int(w * scale_factor)
resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
return resized
elif img.ndim == 3:
c, h, w = img.shape
new_h, new_w = int(h * scale_factor), int(w * scale_factor)
img_hw_c = img.transpose(1, 2, 0)
resized = cv2.resize(img_hw_c, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
return resized.transpose(2, 0, 1)
else:
raise ValueError("Input array must be 2D or 3D")
def normalize_img(img):
if img.dtype == np.uint8:
return img.astype(np.float32) / 255.0, 255
elif img.dtype == np.uint16:
return img.astype(np.float32) / 65535.0, 65535
else:
return img.astype(np.float32), 255
def calc_rmse(pred, target):
"""计算RMSE，输入均为numpy整数数组"""
pred = pred.astype(np.float64)
target = target.astype(np.float64)
return np.sqrt(np.mean((pred - target) ** 2))
def main():
scale = 4
model = AHMF(scale=scale, act='PReLU', in_channels=1, guidance_channels=3)
ckpt_path = f'/home/jianruobing/code/AHMF-main/AHMF-test/model_x{scale}.pth'
checkpoint = torch.load(ckpt_path, map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
model = model.cuda().eval()
input_dir = '/home/jianruobing/code/AHMF-main/imgs'
output_dir = './output'
os.makedirs(output_dir, exist_ok=True)
exts = ['png', 'jpg', 'jpeg']
all_files = []
for ext in exts:
all_files.extend(glob.glob(os.path.join(input_dir, f'*_depth.{ext}')))
all_files.extend(glob.glob(os.path.join(input_dir, f'*_depth_x4.{ext}')))
all_files.extend(glob.glob(os.path.join(input_dir, f'*_rgb.{ext}')))
groups = {}
for fpath in all_files:
fname = os.path.basename(fpath)
if '_depth_x4' in fname:
prefix = fname.split('_depth_x4')[0]
key = 'depth_x4'
elif '_depth' in fname:
prefix = fname.split('_depth')[0]
key = 'depth'
elif '_rgb' in fname:
prefix = fname.split('_rgb')[0]
key = 'rgb'
else:
continue
groups.setdefault(prefix, {})[key] = fpath
print(f'Found {len(groups)} image groups.')
for prefix, files in groups.items():
if not {'depth', 'depth_x4', 'rgb'}.issubset(files.keys()):
print(f'Warning: missing files for {prefix}, skipped.')
continue
dep = cv2.imread(files['depth'], cv2.IMREAD_UNCHANGED)
lr = cv2.imread(files['depth_x4'], cv2.IMREAD_UNCHANGED)
rgb = cv2.imread(files['rgb'], cv2.IMREAD_COLOR)
print(f"{prefix}: Input dep dtype = {dep.dtype}, shape = {dep.shape}")
print(f"{prefix}: Input rgb dtype = {rgb.dtype}, shape = {rgb.shape}")
print(f"{prefix}: Input lr dtype = {lr.dtype}, shape = {lr.shape}")
if dep is None or lr is None or rgb is None:
print(f'Failed loading images for {prefix}')
continue
rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
# 归一化输入图像，获取最大值（255或65535）
dep_norm, dep_max = normalize_img(dep)
lr_norm, lr_max = normalize_img(lr)
rgb_norm = rgb.astype(np.float32) / 255.0
dep_norm, rgb_norm = mod_crop(np.expand_dims(dep_norm, 0), rgb_norm.transpose(2, 0, 1), modulo=32)
lr_norm = mod_crop(lr_norm, modulo=32)
# 上采样 lr
lr_up = imresize(lr_norm, scale)
if lr_norm.ndim == 2:
lr_norm = np.expand_dims(lr_norm, 0)
if lr_up.ndim == 2:
lr_up = np.expand_dims(lr_up, 0)
dep_t, rgb_t, lr_t, lr_up_t = np_to_tensor(dep_norm, rgb_norm, lr_norm, lr_up)
dep_t = dep_t.unsqueeze(0).cuda()
rgb_t = rgb_t.unsqueeze(0).cuda()
lr_t = lr_t.unsqueeze(0).cuda()
lr_up_t = lr_up_t.unsqueeze(0).cuda()
inputs = {
'img_gt': dep_t,
'img_rgb': rgb_t,
'img_lr': lr_t,
'lr_up': lr_up_t,
}
with torch.no_grad():
out = model(inputs)['img_out']
# 反归一化到原始像素范围
out_img = out.squeeze().cpu().numpy()
out_img = np.clip(out_img, 0, 1) * dep_max
# 转为对应整数格式
if dep.dtype == np.uint8:
out_img_int = np.round(out_img).astype(np.uint8)
dep_int = dep
elif dep.dtype == np.uint16:
out_img_int = np.round(out_img).astype(np.uint16)
dep_int = dep
else:
out_img_int = out_img.astype(np.float32)
dep_int = dep.astype(np.float32)
# 计算RMSE，基于整型图像
rmse_val = calc_rmse(out_img_int, dep_int)
print(f'{prefix}: RMSE = {rmse_val:.4f}')
save_path = os.path.join(output_dir, f'{prefix}_sr.png')
# 保存时强制转换为兼容格式
cv2.imwrite(save_path, out_img_int)
print(f'Saved SR image to {save_path}')
if __name__ == '__main__':
main()ASSISTANT
你给的代码主要是用cv2读取图像的，同时你想“参照下面代码中读取npy的方式”，看起来是想把图像读取与预处理方式和`quick_test.py`里对npy读取及处理的兼容性结合起来。
在你现有代码中，读取深度图、低分辨率图为`cv2.imread(..., cv2.IMREAD_UNCHANGED)`，RGB图为彩色读取（然后转RGB），然后对深度图进行归一化（`normalize_img`），对RGB做归一化后转CHW，`mod_crop`裁剪后转tensor，送入模型。
而`quick_test.py`里是直接用`np.load`加载npy数组，然后类似归一化和裁边。
---
### 你要“参照读取npy的方式”，我理解的主要改动点是：
- 你要读取的`depth`、`depth_x4`、`rgb`都是用npy格式存在的。
- 你要对npy数组进行同样的归一化、裁剪及tensor转化处理。
- 后续模型输入和RMSE计算逻辑保持不变。
---
### 基于你给的代码，我给你一份修改方案：
```python
import os
import glob
import torch
import numpy as np
from ahmf import AHMF
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
def mod_crop(*imgs, modulo):
def crop(img):
if img.ndim == 2:
h, w = img.shape
crop_h, crop_w = h % modulo, w % modulo
return img[crop_h // 2 : h - (crop_h - crop_h // 2),
crop_w // 2 : w - (crop_w - crop_w // 2)]
elif img.ndim == 3:
c, h, w = img.shape
crop_h, crop_w = h % modulo, w % modulo
return img[:, crop_h // 2 : h - (crop_h - crop_h // 2),
crop_w // 2 : w - (crop_w - crop_w // 2)]
else:
raise ValueError(f"Unsupported image dimension: {img.ndim}")
out = [crop(img) for img in imgs]
return out if len(out) > 1 else out[0]
def np_to_tensor(*imgs, input_range=1.0, output_range=1.0):
def _to_tensor(img):
arr = img.astype(np.float32)
if arr.ndim == 2:
arr = np.expand_dims(arr, 0)
tensor = torch.from_numpy(arr).float()
tensor *= output_range / input_range
return tensor
out = [_to_tensor(img) for img in imgs]
return out if len(out) > 1 else out[0]
def normalize_img(img):
# 保持原数据类型，并归一化（和quick_test类似）
if img.dtype == np.uint8:
return img.astype(np.float32) / 255.0, 255
elif img.dtype == np.uint16:
return img.astype(np.float32) / 65535.0, 65535
else:
# float类型或者其它
max_val = float(np.max(img)) if np.max(img) > 1.0 else 1.0
return img.astype(np.float32) / max_val, max_val
def calc_rmse(pred, target):
pred = pred.astype(np.float64)
target = target.astype(np.float64)
return np.sqrt(np.mean((pred - target) ** 2))
def main():
scale = 4
model = AHMF(scale=scale, act='PReLU', in_channels=1, guidance_channels=3)
ckpt_path = f'/home/jianruobing/code/AHMF-main/AHMF-test/model_x{scale}.pth'
checkpoint = torch.load(ckpt_path, map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
model = model.cuda().eval()
input_dir = '/home/jianruobing/code/AHMF-main/imgs'
output_dir = './output'
os.makedirs(output_dir, exist_ok=True)
# 查找npy文件
all_files = []
all_files.extend(glob.glob(os.path.join(input_dir, '*_depth.npy')))
all_files.extend(glob.glob(os.path.join(input_dir, '*_depth_x4.npy')))
all_files.extend(glob.glob(os.path.join(input_dir, '*_rgb.npy')))
# 按前缀分组
groups = {}
for fpath in all_files:
fname = os.path.basename(fpath)
if '_depth_x4' in fname:
prefix = fname.split('_depth_x4')[0]
key = 'depth_x4'
elif '_depth' in fname:
prefix = fname.split('_depth')[0]
key = 'depth'
elif '_rgb' in fname:
prefix = fname.split('_rgb')[0]
key = 'rgb'
else:
continue
groups.setdefault(prefix, {})[key] = fpath
print(f'Found {len(groups)} image groups.')
for prefix, files in groups.items():
if not {'depth', 'depth_x4', 'rgb'}.issubset(files.keys()):
print(f'Warning: missing files for {prefix}, skipped.')
continue
dep = np.load(files['depth'])
lr = np.load(files['depth_x4'])
rgb = np.load(files['rgb'])
print(f"{prefix}: Input dep dtype = {dep.dtype}, shape = {dep.shape}")
print(f"{prefix}: Input rgb dtype = {rgb.dtype}, shape = {rgb.shape}")
print(f"{prefix}: Input lr dtype = {lr.dtype}, shape = {lr.shape}")
# 归一化
dep_norm, dep_max = normalize_img(dep)
lr_norm, lr_max = normalize_img(lr)
# rgb假设shape H, W, C，转CHW并归一化（0~1）
rgb_norm = np.float32(np.transpose(rgb, (2, 0, 1))) / 255.0
# mod_crop，深度和rgb都扩1维（channel），rgb转CHW在前面完成了
dep_norm = mod_crop(np.expand_dims(dep_norm, 0), modulo=32)
rgb_norm = mod_crop(rgb_norm, modulo=32)
lr_norm = mod_crop(lr_norm, modulo=32)
# lr上采样用cv2一样双线性插值
h, w = lr_norm.shape[-2], lr_norm.shape[-1] # 对应深度h,w
import cv2
lr_up = cv2.resize(lr_norm, (w * scale, h * scale), interpolation=cv2.INTER_LINEAR)
# 确保shape维度对应
if lr_norm.ndim == 2:
lr_norm = np.expand_dims(lr_norm, 0)
if lr_up.ndim == 2:
lr_up = np.expand_dims(lr_up, 0)
# 转tensor并增加batch维
dep_t, rgb_t, lr_t, lr_up_t = np_to_tensor(dep_norm, rgb_norm, lr_norm, lr_up)
dep_t = dep_t.unsqueeze(0).cuda()
rgb_t = rgb_t.unsqueeze(0).cuda()
lr_t = lr_t.unsqueeze(0).cuda()
lr_up_t = lr_up_t.unsqueeze(0).cuda()
inputs = {
'img_gt': dep_t,
'img_rgb': rgb_t,
'img_lr': lr_t,
'lr_up': lr_up_t,
}
with torch.no_grad():
out = model(inputs)['img_out']
# 反归一化还原
out_img = out.squeeze().cpu().numpy()
out_img = np.clip(out_img, 0, 1) * dep_max
# 转回整数格式做RMSE计算
if dep.dtype == np.uint8:
out_img_int = np.round(out_img).astype(np.uint8)
dep_int = dep
elif dep.dtype == np.uint16:
out_img_int = np.round(out_img).astype(np.uint16)
dep_int = dep
else:
# float或者其它类型（如int32），保持float32，RMSE时会自动转换
out_img_int = out_img.astype(np.float32)
dep_int = dep.astype(np.float32)
rmse_val = calc_rmse(out_img_int, dep_int)
print(f'{prefix}: RMSE = {rmse_val:.4f}')
# 保存结果图为png，需先转uint8或uint16
save_path = os.path.join(output_dir, f'{prefix}_sr.png')
import cv2
cv2.imwrite(save_path, out_img_int)
print(f'Saved SR image to {save_path}')
if __name__ == '__main__':
main()
```
---
### 说明
- 读取`npy`方式用的是`np.load`，对应你给的`quick_test.py`读取方式。
- 归一化方式和`quick_test.py`的处理保持吻合（归一化到[0,1]，保留最大值）。
- rgb从HWC转换为CHW：`np.transpose(rgb, (2,0,1))`，并除255归一化。
- `mod_crop`改造成接受CHW格式图像。
- 双线性插值用cv2的`cv2.resize`实现，与之前函数一致。
- 转Tensor部分用你示例里的`np_to_tensor`。
- RMSE计算调用你的`calc_rmse`。
- 输出后恢复到原始像素值范围，转为整数类型保存PNG。
这样修改后，你的输入数据结构和预处理流程基本与`quick_test.py`似的，方便对接。
如果你还有其它细节需要同步，可以告诉我！