From fd6671d3d4f31aad439df349a9406465729e57c5 Mon Sep 17 00:00:00 2001 From: "li.zhenye" Date: Mon, 8 Aug 2022 11:27:38 +0800 Subject: [PATCH] =?UTF-8?q?[ext]=20=E6=B7=BB=E5=8A=A0ENVI=E5=9B=BE?= =?UTF-8?q?=E7=89=87=E8=BD=AC=E6=8D=A2=E5=8A=9F=E8=83=BDBeta?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 2 +- main_test.py | 136 ++++++++++++++++++++++++++++++++++++--------------- utils.py | 20 ++++++++ 3 files changed, 117 insertions(+), 41 deletions(-) diff --git a/main.py b/main.py index 65c5fea..7a63381 100755 --- a/main.py +++ b/main.py @@ -86,7 +86,7 @@ def main(only_spec=False, only_color=False): # 进行多个喷阀的合并 masks = [utils.valve_expend(mask) for mask in [mask_spec, mask_rgb]] # 进行喷阀同时开启限制 - + masks = [utils.valve_limit(mask, Config.max_open_valve_limit) for mask in masks] # control the size of the output masks, 在resize前，图像的宽度是和喷阀对应的 masks = [cv2.resize(mask.astype(np.uint8), Config.target_size) for mask in masks] # 写出 diff --git a/main_test.py b/main_test.py index 0990385..2c00bca 100644 --- a/main_test.py +++ b/main_test.py @@ -3,18 +3,22 @@ # @Auther : zhouchao # @File: main_test.py # @Software:PyCharm +import datetime import itertools import logging import os import time import socket import typing +from shutil import copyfile import cv2 import matplotlib.pyplot as plt import numpy as np +from matplotlib.gridspec import GridSpec import transmit +import utils from config import Config from models import Detector, AnonymousColorDetector, ManualTree, SpecDetector, RgbDetector from utils import read_labeled_img, size_threshold, natural_sort @@ -28,14 +32,16 @@ class TestMain: background_model_path=Config.rgb_background_model_path) def pony_run(self, test_path, test_spectra=False, test_rgb=False, - convert=False, get_delta=False): + convert_dir=None, get_delta=False, silent=False): """ 虚拟读图测试程序 :param test_path: 测试文件夹或者图片 :param test_spectra: 是否测试光谱 :param test_rgb: 是否测试rgb - :param convert: 是否进行格式转化 + :param convert_dir: 格式转化文件夹 + :param get_delta:是否计算偏差 + :param silent: 是否静默模式运行 :return: """ if os.path.isdir(test_path): @@ -55,24 +61,40 @@ class TestMain: _ = self.test_rgb(rgb_img, img_name=test_path) return 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[:] + assert rgb_temp.replace('rgb', '') == spec_temp.replace('spec', '') if test_spectra: + print(f"test spec file {spec_file_name}") with open(os.path.join(test_path, spec_file_name), 'rb') as f: data = f.read() - spec_img = transmit.BeforeAfterMethods.spec_data_post_process(data) - if convert: - # TODO: 完善这个文件转换功能 + try: + spec_img = transmit.BeforeAfterMethods.spec_data_post_process(data) + except Exception as e: + raise FileExistsError(f'文件 {spec_file_name} 读取失败. 请清理文件夹{test_path},去除{spec_file_name}') + + 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), dtype=np.uint8) - cv2.imwrite(rgb_file_name + '.bmp', spec_img_show[..., ::-1]) + cv2.imwrite(os.path.join(convert_dir, spec_file_name + '.bmp'), spec_img_show[..., ::-1]) + hdr_string = utils.generate_hdr(f'File imported from {spec_file_name} at' + f' {datetime.datetime.now().strftime("%m/%d/%Y, %H:%M")}') + copyfile(os.path.join(test_path, spec_file_name), os.path.join(convert_dir, spec_file_name+'.raw')) - spec_mask = self.test_spec(spec_img, img_name=spec_file_name) + with open(os.path.join(convert_dir, spec_file_name + '.hdr'), 'w') as f: + f.write(hdr_string) + spec_mask = self.test_spec(spec_img, img_name=spec_file_name, show=False) if test_rgb: + print(f'test rgb file {rgb_file_name}') with open(os.path.join(test_path, rgb_file_name), 'rb') as f: data = f.read() - rgb_img = transmit.BeforeAfterMethods.rgb_data_post_process(data) - if convert: - cv2.imwrite(os.path.join(test_path, rgb_file_name + '.bmp'), rgb_img[..., ::-1]) - rgb_mask = self.test_rgb(rgb_img, img_name=rgb_file_name) + try: + rgb_img = transmit.BeforeAfterMethods.rgb_data_post_process(data) + except Exception as e: + raise FileExistsError(f'文件 {rgb_file_name} 读取失败. 请清理文件夹{test_path}, 去除{rgb_file_name}') + + if convert_dir is not None: + cv2.imwrite(os.path.join(convert_dir, rgb_file_name + '.bmp'), rgb_img[..., ::-1]) + rgb_mask = self.test_rgb(rgb_img, img_name=rgb_file_name, show=False) if test_rgb and test_spectra: if get_delta: spec_cv = np.clip(spec_img[..., [21, 3, 0]], a_min=0, a_max=1) * 255 @@ -81,42 +103,58 @@ class TestMain: print(delta) self.merge(rgb_img=rgb_img, rgb_mask=rgb_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) - def test_rgb(self, rgb_img, img_name): + def test_rgb(self, rgb_img, img_name, show=True): rgb_mask = self._rgb_detector.predict(rgb_img) - fig, axs = plt.subplots(2, 1) - axs[0].imshow(rgb_img) - axs[0].set_title(f"rgb img {img_name}") - axs[1].imshow(rgb_mask) - axs[1].set_title('rgb mask') - plt.show() + if show: + fig, axs = plt.subplots(2, 1) + axs[0].imshow(rgb_img) + axs[0].set_title(f"rgb img {img_name}") + axs[1].imshow(rgb_mask) + axs[1].set_title('rgb mask') + plt.show() return rgb_mask - def test_spec(self, spec_img, img_name): + def test_spec(self, spec_img, img_name, show=True): spec_mask = self._spec_detector.predict(spec_img) - fig, axs = plt.subplots(2, 1) - axs[0].imshow(spec_img[..., [21, 3, 0]]) - axs[0].set_title(f"spec img {img_name}") - axs[1].imshow(spec_mask) - axs[1].set_title('spec mask') - plt.show() + if show: + fig, axs = plt.subplots(2, 1) + axs[0].imshow(spec_img[..., [21, 3, 0]]) + axs[0].set_title(f"spec img {img_name}") + axs[1].imshow(spec_mask) + axs[1].set_title('spec mask') + plt.show() return spec_mask @staticmethod - def merge(rgb_img, rgb_mask, spec_img, spec_mask, rgb_file_name, spec_file_name): + def merge(rgb_img, rgb_mask, spec_img, spec_mask, rgb_file_name, spec_file_name, show=True): mask_result = (spec_mask | rgb_mask).astype(np.uint8) mask_result = mask_result.repeat(Config.blk_size, axis=0).repeat(Config.blk_size, axis=1).astype(np.uint8) - fig, axs = plt.subplots(3, 2) - axs[0, 0].set_title(rgb_file_name) - axs[0, 0].imshow(rgb_img) - axs[1, 0].imshow(spec_img) - axs[1, 0].set_title(spec_file_name) - axs[2, 0].imshow(mask_result) - axs[0, 1].imshow(rgb_mask) - axs[1, 1].imshow(spec_mask) - axs[2, 1].imshow(mask_result) - plt.show() + if show: + fig = plt.figure(constrained_layout=True) + gs = GridSpec(3, 2, figure=fig) + ax1 = fig.add_subplot(gs[0, 0]) + ax2 = fig.add_subplot(gs[0, 1]) + ax3 = fig.add_subplot(gs[1, 0]) + ax4 = fig.add_subplot(gs[1, 1]) + ax5 = fig.add_subplot(gs[2, :]) + fig.suptitle("Result") + + ax1.imshow(rgb_img) + ax1.set_title(rgb_file_name) + ax2.imshow(rgb_mask) + ax2.set_title("rgb mask") + + ax3.imshow(np.clip(spec_img, a_min=0, a_max=1)) + ax3.set_title(spec_file_name) + ax4.imshow(spec_mask) + ax4.set_title('spec mask') + + ax5.imshow(mask_result) + ax5.set_title('merge mask') + plt.show() return mask_result def calculate_delta(self, rgb_img, spec_img, search_area_size=(400, 200), eps=1): @@ -173,8 +211,26 @@ class TestMain: if __name__ == '__main__': import argparse + parser = argparse.ArgumentParser(description='Run image test or image_file') + parser.add_argument('path', default=r'E:\zhouchao\8.4\zazhi', help='测试文件或文件夹') + parser.add_argument('-test_rgb', default=True, action='store_false', help='是否测试RGB图') + parser.add_argument('-test_spec', default=True, action='store_false', help='是否测试光谱图') + parser.add_argument('-get_delta', default=False, action='store_true', help='是否进行误差计算') + parser.add_argument('-convert_dir', default=None, help='是否将c语言采集的buffer进行转换') + parser.add_argument('-s', '--silent', default=False, action='store_true', help='是否显示') + args = parser.parse_args() + # file check + if args.convert_dir is not None: + if os.path.exists(args.convert_dir): + if not os.path.isdir(args.convert_dir): + raise TypeError("转换文件夹需要是个文件夹") + else: + if os.path.abspath(args.path) == os.path.abspath(args.convert_dir): + print("警告！您的输出文件夹和输入文件夹用了同一个位置") + else: + print(f"已创建需要存放转换文件的文件夹 {args.convert_dir}") + os.makedirs(args.convert_dir, mode=0o777, exist_ok=False) - parser = argparse.ArgumentParser(description='Run image test or ') tester = TestMain() - tester.pony_run(test_path=r'E:\zhouchao\8.4\zazhi', - test_rgb=True, test_spectra=True, get_delta=False, convert=False) + tester.pony_run(test_path=args.path, test_rgb=args.test_rgb, test_spectra=args.test_spec, + get_delta=args.get_delta, convert_dir=args.convert_dir, silent=args.silent) diff --git a/utils.py b/utils.py index 9e6c4ea..e4b9cf4 100755 --- a/utils.py +++ b/utils.py @@ -14,6 +14,8 @@ from numpy.random import default_rng from matplotlib import pyplot as plt import re +from config import Config + def natural_sort(l): """ @@ -254,6 +256,24 @@ def read_envi_ascii(file_name, save_xy=False, hdr_file_name=None): return dict(zip(class_names, vectors)) +def generate_hdr(des='File Imported into ENVI.'): + template_file = f"""ENVI +description = {{ {des} + }} +samples = {Config.nCols} +lines = {Config.nRows} +bands = {Config.nBands} +header offset = 0 +file type = ENVI Standard +data type = 4 +interleave = bil +sensor type = Unknown +byte order = 0 +wavelength units = Unknown +""" + return template_file + + if __name__ == '__main__': # color_dict = {(0, 0, 255): "yangeng", (255, 0, 0): "bejing", (0, 255, 0): "hongdianxian", # (255, 0, 255): "chengsebangbangtang", (0, 255, 255): "lvdianxian"}