def test_model(args):
net = ISONet()
if torch.cuda.is_available():
device = torch.cuda.current_device()
else:
device = 'cpu'
net = net.to(device=device)
model_path = args.model_path
checkpoint = args.checkpoint
# Load the model
print(f"load model:{model_path}")
if checkpoint:
net.load_state_dict(
torch.load(model_path, map_location=torch.device(device))["net"])
else:
net.load_state_dict(
torch.load(model_path, map_location=torch.device(device)))
# Load the test image,BGR -> RGB
img = cv2.imread(args.pic_path)[:, :, ::-1]
# HWC ->CHW
img = img.transpose(2, 0, 1)
# normalize
img = normalize(img)
# the test image is split to 64*64 patches
patches = img_2_patches(img, 64, 64)
c, h, w = img.shape
Hb = int(np.floor(h / 64))
Wb = int(np.floor(w / 64))
print(f"patches {patches.shape[3]}")
x = np.linspace(1, patches.shape[3], patches.shape[3])
y = []
res = []
start_time = time.time()
for nx in range(patches.shape[3]):
with torch.no_grad():
p = torch.from_numpy(patches[:, :, :,
nx]).to(dtype=torch.float32,
device=device).unsqueeze(0)
pre = net(p)
value = pre.item()
res.append(value)
y.append(value)
y = np.array(y)
end_time = time.time()
print(f"Running time:{(end_time-start_time):.2f}s")
# (2) of the paper
print(f"Estimated ISO metric: {np.median(y):.3f}")
# plot
res = np.array(res)
plt.imshow(res.reshape([Hb, Wb]))
plt.colorbar()
plt.show()
def gen_data(args):
"""
generate datasets for training and validation
"""
data_path = args.data_path
save_path = args.save_path
train = args.train
test = args.test
size = args.size
stride = args.stride
aug_times = args.aug_times
gray_mode = args.gray_mode
pic_type = args.pic_type
train_path = Path(data_path).joinpath("Train")
val_data_path = Path(data_path).joinpath("Test")
if save_path is not None:
save_path = Path(save_path)
if not save_path.exists():
save_path.mkdir()
files_train = {}
files_test = {}
for x in train_path.glob("*"):
if x.is_dir():
file_list_train = [str(f_train.absolute().resolve()) for f_train in x.glob(f"*.{pic_type}")]
files_train[x.name] = []
files_train[x.name].extend(file_list_train)
for y in val_data_path.glob("*"):
if y.is_dir():
file_list_test = [str(f_test.absolute().resolve()) for f_test in y.glob(f"*.{pic_type}")]
files_test[y.name] = []
files_test[y.name].extend(file_list_test)
if gray_mode:
train_h5 = 'train_gray.h5'
train_h5_label = 'train_gray_label.h5'
val_h5 = 'val_gray.h5'
val_h5_label = 'val_gray_label.h5'
else:
train_h5 = 'train_rgb.h5'
train_h5_label = 'train_rgb_label.h5'
val_h5 = 'val_rgb.h5'
val_h5_label = 'val_rgb_label.h5'
if train:
# load the training img and generate the dataset
f_train = h5py.File(save_path.joinpath(train_h5), 'w')
f_train_label = h5py.File(save_path.joinpath(train_h5_label), 'w')
train_num = 0
# k->label v->filename list
for k, v in files_train.items():
print(k)
print(v)
if len(v) == 0:
continue
# load a full size image
for f in v:
if gray_mode:
# H * W * C
t_pic = cv2.imread(f, cv2.IMREAD_GRAYSCALE)
else:
t_pic = cv2.imread(f, cv2.IMREAD_COLOR)
# BRG -> RGB
t_pic = t_pic[:, :, ::-1]
# HWC -> CHW
t_pic = np.transpose(t_pic, (2, 0, 1))
t_pic = normalize(t_pic)
# CHW * patch_size
patches = img_2_patches(t_pic, size, stride)
# Control the maximum sample from a single image
patches = patches[:, :, :, :2400]
# dealing with every patch
print(f"training file:{f} --> ##{patches.shape[3]}##sample")
for nx in range(patches.shape[3]):
data = patches[:, :, :, nx]
f_train.create_dataset(str(train_num), data=data)
f_train_label.create_dataset(str(train_num), data=np.array(get_label(int(k))))
train_num += 1
# data augmentation
for mx in range(aug_times):
data_aug = data_augmentation(patches[:, :, :, nx].copy(), np.random.randint(1, 8))
f_train.create_dataset(str(train_num), data=data_aug)
f_train_label.create_dataset(str(train_num), data=np.array(get_label(int(k))))
train_num += 1
f_train.close()
f_train_label.close()
print(f"the number of training images:{train_num}")
if test:
# Gen Test Data
f_test = h5py.File(save_path.joinpath(val_h5), 'w')
f_test_label = h5py.File(save_path.joinpath(val_h5_label), 'w')
# k->label v->filename list
val_num = 0
for k, v in files_test.items():
print(k)
print(v)
if len(v) == 0:
continue
# load full size image
for f in v:
if gray_mode:
# H * W * C
t_pic = cv2.imread(f, cv2.IMREAD_GRAYSCALE)
else:
t_pic = cv2.imread(f, cv2.IMREAD_COLOR)
# BRG -> RGB
t_pic = t_pic[:, :, ::-1]
# HWC -> CHW
t_pic = np.transpose(t_pic, (2, 0, 1))
t_pic = normalize(t_pic)
# CHW * patch_size
patches = img_2_patches(t_pic, size, stride)
# dealing with every patch
print(f"validation file:{f} --> ##{patches.shape[3]}##sample")
for nx in range(patches.shape[3]):
data = patches[:, :, :, nx]
f_test.create_dataset(str(val_num), data=data)
f_test_label.create_dataset(str(val_num), data=np.array(get_label(int(k))))
val_num += 1
f_test.close()
f_test_label.close()
print(f"the number of validation images:{val_num}")
def test_model(args):
net = ISONet()
if torch.cuda.is_available():
device = torch.cuda.current_device()
else:
device = 'cpu'
net = net.to(device=device)
model_path = args.model_path
checkpoint = args.checkpoint
print(f"load model:{model_path}")
if checkpoint:
# 加载模型
net.load_state_dict(
torch.load(model_path, map_location=torch.device(device))["net"])
else:
net.load_state_dict(
torch.load(model_path, map_location=torch.device(device)))
# 读入测试图片,BGR -> RGB
img = cv2.imread(args.pic_path)[:, :, ::-1]
# HWC ->CHW
img = img.transpose(2, 0, 1)
# 归一化
img = normalize(img)
# 把图片分割为若干小块
patches = img_2_patches(img, 64, 64)
c, h, w = img.shape
Hb = int(np.floor(h / 64))
Wb = int(np.floor(w / 64))
print(f"patches {patches.shape[3]}")
x = np.linspace(1, patches.shape[3], patches.shape[3])
y = []
res = []
start_time = time.time()
for nx in range(patches.shape[3]):
with torch.no_grad():
p = torch.from_numpy(patches[:, :, :,
nx]).to(dtype=torch.float32,
device=device).unsqueeze(0)
pre = net(p)
value = pre.item()
res.append(value)
y.append(value)
predict_iso = math.pow(2, pre.item()) * 100
# print(predict_iso)
y = np.array(y)
end_time = time.time()
print(f"使用时间：{(end_time-start_time):.2f}s")
# plot scatter
plt.scatter(x, y)
plt.xlabel('index')
plt.ylabel('predict_iso')
plt.show()
# plot
res = np.array(res)
plt.imshow(res.reshape([Hb, Wb]))
plt.colorbar()
plt.show()
def gen_data(args):
"""
生成训练测试数据集
"""
data_path = args.data_path
save_path = args.save_path
train = args.train
test = args.test
size = args.size
stride = args.stride
aug_times = args.aug_times
gray_mode = args.gray_mode
pic_type = args.pic_type
train_path = Path(data_path).joinpath("Train")
val_data_path = Path(data_path).joinpath("Test")
if save_path is not None:
save_path = Path(save_path)
if not save_path.exists():
save_path.mkdir()
files_train = {}
files_test = {}
for x in train_path.glob("*"):
if x.is_dir():
file_list_train = [
str(f_train.absolute().resolve())
for f_train in x.glob(f"*.{pic_type}")
]
files_train[x.name] = []
files_train[x.name].extend(file_list_train)
for y in val_data_path.glob("*"):
if y.is_dir():
file_list_test = [
str(f_test.absolute().resolve())
for f_test in y.glob(f"*.{pic_type}")
]
files_test[y.name] = []
files_test[y.name].extend(file_list_test)
if gray_mode:
train_h5 = 'train_gray.h5'
train_h5_label = 'train_gray_label.h5'
val_h5 = 'val_gray.h5'
val_h5_label = 'val_gray_label.h5'
else:
train_h5 = 'train_rgb.h5'
train_h5_label = 'train_rgb_label.h5'
val_h5 = 'val_rgb.h5'
val_h5_label = 'val_rgb_label.h5'
if train:
# 读取训练图片，并成生数据集
f_train = h5py.File(save_path.joinpath(train_h5), 'w')
f_train_label = h5py.File(save_path.joinpath(train_h5_label), 'w')
train_num = 0
# k->label v->filename list
for k, v in files_train.items():
print(k)
print(v)
if len(v) == 0:
continue
# 读取每一张大图
for f in v:
if gray_mode:
# H * W * C
t_pic = cv2.imread(f, cv2.IMREAD_GRAYSCALE)
else:
t_pic = cv2.imread(f, cv2.IMREAD_COLOR)
# BRG -> RGB
t_pic = t_pic[:, :, ::-1]
# HWC -> CHW
t_pic = np.transpose(t_pic, (2, 0, 1))
t_pic = normalize(t_pic)
# CHW * patch_size
patches = img_2_patches(t_pic, size, stride)
# 控制样本数量
patches = patches[:, :, :, :2400]
# 处理每一张小图
print(f"训练文件:{f} --> ##{patches.shape[3]}##样本")
for nx in range(patches.shape[3]):
data = patches[:, :, :, nx]
f_train.create_dataset(str(train_num), data=data)
f_train_label.create_dataset(str(train_num),
data=np.array(
get_label(int(k))))
train_num += 1
# 数据增广
for mx in range(aug_times):
data_aug = data_augmentation(
patches[:, :, :, nx].copy(),
np.random.randint(1, 8))
f_train.create_dataset(str(train_num), data=data_aug)
f_train_label.create_dataset(str(train_num),
data=np.array(
get_label(int(k))))
train_num += 1
f_train.close()
f_train_label.close()
print(f"训练集图片数量:{train_num}")
if test:
# Gen Test Data
f_test = h5py.File(save_path.joinpath(val_h5), 'w')
f_test_label = h5py.File(save_path.joinpath(val_h5_label), 'w')
# k->label v->filename list
val_num = 0
for k, v in files_test.items():
print(k)
print(v)
if len(v) == 0:
continue
# 读取每一张大图
for f in v:
if gray_mode:
# H * W * C
t_pic = cv2.imread(f, cv2.IMREAD_GRAYSCALE)
else:
t_pic = cv2.imread(f, cv2.IMREAD_COLOR)
# BRG -> RGB
t_pic = t_pic[:, :, ::-1]
# HWC -> CHW
t_pic = np.transpose(t_pic, (2, 0, 1))
t_pic = normalize(t_pic)
# CHW * patch_size
patches = img_2_patches(t_pic, size, stride)
# 处理每一张小图
print(f"测试文件:{f} --> ##{patches.shape[3]}##样本")
for nx in range(patches.shape[3]):
data = patches[:, :, :, nx]
f_test.create_dataset(str(val_num), data=data)
f_test_label.create_dataset(str(val_num),
data=np.array(get_label(
int(k))))
val_num += 1
f_test.close()
f_test_label.close()
print(f"测试集图片数量:{val_num}")