NanjingForestryUniversity/supermachine-tobacco
3
0
Fork 0
mirror of https://github.com/NanjingForestryUniversity/supermachine-tobacco.git synced 2025-11-08 14:23:55 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

[ext] 增加了格式转换时的进度条

Browse Source
This commit is contained in:
li.zhenye 2022-08-09 11:39:11 +08:00
parent 483a78d423
commit 3cdafcd6fd
1 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
main_test.py
View File

@ -16,6 +16,7 @@ import cv2
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import numpy as np import numpy as np
from matplotlib.gridspec import GridSpec from matplotlib.gridspec import GridSpec
import tqdm
import transmit import transmit
import utils import utils
@ -60,17 +61,19 @@ class TestMain:
rgb_img = transmit.BeforeAfterMethods.rgb_data_post_process(data) rgb_img = transmit.BeforeAfterMethods.rgb_data_post_process(data)
_ = self.test_rgb(rgb_img, img_name=test_path) _ = self.test_rgb(rgb_img, img_name=test_path)
return return
if silent and convert_dir is not None:
bar = tqdm.tqdm(desc='正在转换', total=len(rgb_file_names), ncols=80)
for rgb_file_name, spec_file_name in zip(rgb_file_names, spec_file_names): for rgb_file_name, spec_file_name in zip(rgb_file_names, spec_file_names):
rgb_temp, spec_temp = rgb_file_name[:], spec_file_name[:] rgb_temp, spec_temp = rgb_file_name[:], spec_file_name[:]
assert rgb_temp.replace('rgb', '') == spec_temp.replace('spec', '') assert rgb_temp.replace('rgb', '') == spec_temp.replace('spec', '')
if test_spectra: if test_spectra:
print(f"test spec file {spec_file_name}")
with open(os.path.join(test_path, spec_file_name), 'rb') as f: with open(os.path.join(test_path, spec_file_name), 'rb') as f:
data = f.read() data = f.read()
try: try:
spec_img = transmit.BeforeAfterMethods.spec_data_post_process(data) spec_img = transmit.BeforeAfterMethods.spec_data_post_process(data)
except Exception as e: except Exception as e:
raise FileExistsError(f'文件 {spec_file_name} 读取失败. 请清理文件夹{test_path},去除{spec_file_name}') raise FileExistsError(f'文件 {spec_file_name} 读取失败,长度不对. 请清理文件夹{test_path},'
f'去除{spec_file_name}\n {e}')
if convert_dir is not None: if convert_dir is not None:
spec_img_show = np.asarray(np.clip(spec_img[..., [21, 3, 0]] * 255, a_min=0, a_max=255), spec_img_show = np.asarray(np.clip(spec_img[..., [21, 3, 0]] * 255, a_min=0, a_max=255),
@ -84,13 +87,13 @@ class TestMain:
f.write(hdr_string) f.write(hdr_string)
spec_mask = self.test_spec(spec_img, img_name=spec_file_name, show=False) spec_mask = self.test_spec(spec_img, img_name=spec_file_name, show=False)
if test_rgb: if test_rgb:
print(f'test rgb file {rgb_file_name}')
with open(os.path.join(test_path, rgb_file_name), 'rb') as f: with open(os.path.join(test_path, rgb_file_name), 'rb') as f:
data = f.read() data = f.read()
try: try:
rgb_img = transmit.BeforeAfterMethods.rgb_data_post_process(data) rgb_img = transmit.BeforeAfterMethods.rgb_data_post_process(data)
except Exception as e: except Exception as e:
raise FileExistsError(f'文件 {rgb_file_name} 读取失败. 请清理文件夹{test_path}, 去除{rgb_file_name}') raise FileExistsError(f'文件 {rgb_file_name} 读取失败, 长度不对. 请清理文件夹{test_path},'
f' 去除{rgb_file_name} \n {e}')
if convert_dir is not None: if convert_dir is not None:
cv2.imwrite(os.path.join(convert_dir, rgb_file_name + '.bmp'), rgb_img[..., ::-1]) cv2.imwrite(os.path.join(convert_dir, rgb_file_name + '.bmp'), rgb_img[..., ::-1])
@ -105,6 +108,8 @@ class TestMain:
spec_img=spec_img[..., [21, 3, 0]], spec_mask=spec_mask, spec_img=spec_img[..., [21, 3, 0]], spec_mask=spec_mask,
rgb_file_name=rgb_file_name, spec_file_name=spec_file_name, rgb_file_name=rgb_file_name, spec_file_name=spec_file_name,
show=not silent) show=not silent)
if silent and convert_dir is not None:
bar.update()
def test_rgb(self, rgb_img, img_name, show=True): def test_rgb(self, rgb_img, img_name, show=True):
rgb_mask = self._rgb_detector.predict(rgb_img) rgb_mask = self._rgb_detector.predict(rgb_img)