[ext] 添加ENVI图片转换功能Beta

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]
         # 写出

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)
+                try:
-                if convert:
+                    spec_img = transmit.BeforeAfterMethods.spec_data_post_process(data)
-                    # TODO: 完善这个文件转换功能
+                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)
+                try:
-                if convert:
+                    rgb_img = transmit.BeforeAfterMethods.rgb_data_post_process(data)
-                    cv2.imwrite(os.path.join(test_path, rgb_file_name + '.bmp'), rgb_img[..., ::-1])
+                except Exception as e:
-                rgb_mask = self.test_rgb(rgb_img, img_name=rgb_file_name)
+                    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)
+        if show:
-        axs[0].imshow(rgb_img)
+            fig, axs = plt.subplots(2, 1)
-        axs[0].set_title(f"rgb img {img_name}")
+            axs[0].imshow(rgb_img)
-        axs[1].imshow(rgb_mask)
+            axs[0].set_title(f"rgb img {img_name}")
-        axs[1].set_title('rgb mask')
+            axs[1].imshow(rgb_mask)
-        plt.show()
+            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)
+        if show:
-        axs[0].imshow(spec_img[..., [21, 3, 0]])
+            fig, axs = plt.subplots(2, 1)
-        axs[0].set_title(f"spec img {img_name}")
+            axs[0].imshow(spec_img[..., [21, 3, 0]])
-        axs[1].imshow(spec_mask)
+            axs[0].set_title(f"spec img {img_name}")
-        axs[1].set_title('spec mask')
+            axs[1].imshow(spec_mask)
-        plt.show()
+            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)
+        if show:
-        axs[0, 0].set_title(rgb_file_name)
+            fig = plt.figure(constrained_layout=True)
-        axs[0, 0].imshow(rgb_img)
+            gs = GridSpec(3, 2, figure=fig)
-        axs[1, 0].imshow(spec_img)
+            ax1 = fig.add_subplot(gs[0, 0])
-        axs[1, 0].set_title(spec_file_name)
+            ax2 = fig.add_subplot(gs[0, 1])
-        axs[2, 0].imshow(mask_result)
+            ax3 = fig.add_subplot(gs[1, 0])
-        axs[0, 1].imshow(rgb_mask)
+            ax4 = fig.add_subplot(gs[1, 1])
-        axs[1, 1].imshow(spec_mask)
+            ax5 = fig.add_subplot(gs[2, :])
-        axs[2, 1].imshow(mask_result)
+            fig.suptitle("Result")
-        plt.show()
+
+            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',
+    tester.pony_run(test_path=args.path, test_rgb=args.test_rgb, test_spectra=args.test_spec,
-                    test_rgb=True, test_spectra=True, get_delta=False, convert=False)
+                    get_delta=args.get_delta, convert_dir=args.convert_dir, silent=args.silent)

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"}