micro_plastic/only_mask.py

import argparse
from bil2rgb import process_bil_files
from shape_spectral import process_images
import cv2
from classification_model.Parallel.predict_plastic import load_model, predict_with_model
import numpy as np
import os
import matplotlib
from shape_spectral_background import process_images_background
# import pandas as pd
from mask import detect_microplastic_mask_from_array

matplotlib.use('TkAgg')


def read_hdr_file(bil_path):
    hdr_path = bil_path.replace('.bil', '.hdr')
    with open(hdr_path, 'r') as f:
        header = f.readlines()

    samples, lines = None, None

    for line in header:
        if line.startswith('samples'):
            samples = int(line.split('=')[-1].strip())
        if line.startswith('lines'):
            lines = int(line.split('=')[-1].strip())

    return samples, lines


def save_envi_classification(bil_path, df, savepath):
    samples, lines = read_hdr_file(bil_path)
    classification_result = np.zeros((lines, samples), dtype=np.uint16)

    for _, row in df.iterrows():
        contour = row['contour']
        prediction = int(row['Predictions']) + 1
        contour = np.array(contour, dtype=np.int32)
        # 先将 classification_result 中的 10 和 11 替换为 0
        classification_result[(classification_result == 10)] = 0

        cv2.fillPoly(classification_result, [contour], prediction)

    output_path = savepath

    with open(output_path, 'wb') as f:
        classification_result.tofile(f)

    header_content = f"""ENVI
description = {{
Classification Result.}}
samples = {samples}
lines = {lines}
bands = 1
header offset = 0
file type = ENVI Standard
data type = 2
interleave = bil
classes = 11
class = {{ background, ABS, HDPE, LDPE, PA6, PET, PP, PS, PTFE, PVC,background2 }}
single pixel area = 0.000036
unit = mm2
byte order = 0
wavelength units = nm
"""
    filename, ext = os.path.splitext(savepath)
    # 替换扩展名为 '.hdr'
    header_filename = filename + '.hdr'

    with open(header_filename, 'w') as header_file:
        header_file.write(header_content)


def change_hdr_file(bil_path):
    # 定义要追加的波长信息
    wavelength_info = """wavelength = {898.82, 903.64, 908.46, 913.28, 918.1, 922.92, 927.75, 932.57, 937.4, 942.22, 947.05, 951.88, 956.71, 961.54, 966.38, 971.21, 976.05, 980.88, 985.72, 990.56, 995.4, 1000.2, 1005.1, 1009.9, 1014.8, 1019.6, 1024.5, 1029.3, 1034.2, 1039, 1043.9, 1048.7, 1053.6, 1058.4, 1063.3, 1068.2, 1073, 1077.9, 1082.7, 1087.6, 1092.5, 1097.3, 1102.2, 1107.1, 1111.9, 1116.8, 1121.7, 1126.6, 1131.4, 1136.3, 1141.2, 1146.1, 1150.9, 1155.8, 1160.7, 1165.6, 1170.5, 1175.4, 1180.2, 1185.1, 1190, 1194.9, 1199.8, 1204.7, 1209.6, 1214.5, 1219.4, 1224.3, 1229.2, 1234.1, 1239, 1243.9, 1248.8, 1253.7, 1258.6, 1263.5, 1268.4, 1273.3, 1278.2, 1283.1, 1288.1, 1293, 1297.9, 1302.8, 1307.7, 1312.6, 1317.6, 1322.5, 1327.4, 1332.3, 1337.3, 1342.2, 1347.1, 1352, 1357, 1361.9, 1366.8, 1371.8, 1376.7, 1381.6, 1386.6, 1391.5, 1396.5, 1401.4, 1406.3, 1411.3, 1416.2, 1421.2, 1426.1, 1431.1, 1436, 1441, 1445.9, 1450.9, 1455.8, 1460.8, 1465.8, 1470.7, 1475.7, 1480.6, 1485.6, 1490.6, 1495.5, 1500.5, 1505.5, 1510.4, 1515.4, 1520.4, 1525.3, 1530.3, 1535.3, 1540.3, 1545.2, 1550.2, 1555.2, 1560.2, 1565.2, 1570.1, 1575.1, 1580.1, 1585.1, 1590.1, 1595.1, 1600.1, 1605.1, 1610, 1615, 1620, 1625, 1630, 1635, 1640, 1645, 1650, 1655, 1660, 1665, 1670.1, 1675.1, 1680.1, 1685.1, 1690.1, 1695.1, 1700.1, 1705.1, 1710.2, 1715.2, 1720.2}"""

    # 将.bil路径转换为.hdr路径
    hdr_path = os.path.splitext(bil_path)[0] + '.hdr'

    # 检查.hdr文件是否存在
    if not os.path.exists(hdr_path):
        print(f"错误: 找不到对应的HDR文件: {hdr_path}")
        return

    # 读取文件内容
    with open(hdr_path, 'r') as file:
        content = file.read()

    # 检查是否已包含波长信息
    if 'wavelength' in content:
        print(f"File {os.path.basename(hdr_path)} already contains wavelength information; no changes needed.")
        return

    # 检查文件是否以换行符结尾
    needs_newline = not content.endswith('\n')

    # 追加波长信息
    with open(hdr_path, 'a') as file:
        if needs_newline:
            file.write('\n')  # 确保新内容从新行开始
        file.write(wavelength_info + '\n')

    print(f"Successfully added wavelength information to file: {os.path.basename(hdr_path)}")


def main():
    bil_path = "D:\Data\MPData7.bil"


    # 处理BIL文件生成RGB图像
    print("Processing BIL file to generate RGB image...\n")
    rgb_img = process_bil_files(bil_path)

    # 修改hdr
    change_hdr_file(bil_path)

    # 生成掩膜，mask为16位的塑料标签掩膜
    print("Generating mask ...\n")
    mask = detect_microplastic_mask_from_array(
        image=rgb_img,  # 直接传入cv2.imread的结果
        filter_method='threshold',
        diameter=None,
        flow_threshold=0.4,
        cellprob_threshold=-1
    )




if __name__ == "__main__":
    main()