SAR/test_data/xie/mydetection.py

import subprocess
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
from my_config import averagePooling, maxPooling, xyOverlap
import sys
plt.ioff()


def ship_detection_fun(pic_path, label_path):

    sar_img = cv2.imread(pic_path)

    # 读取标签矩阵
    label_matrix = np.loadtxt(label_path)
    label_matrix[:, 2] -= label_matrix[:, 0]
    label_matrix[:, 3] -= label_matrix[:, 1]

    # 转换为灰度图像
    image_init = cv2.cvtColor(sar_img, cv2.COLOR_BGR2GRAY)

    # 进行卷积和池化操作
    image = cv2.filter2D(image_init, -1, np.ones((13, 13)) / 169)
    image = averagePooling(image, [5, 5])

    # 进行 GMM 聚类
    from sklearn.mixture import GaussianMixture

    numComponents = 10
    selected_components = 2
    gmm = GaussianMixture(n_components=numComponents)
    image_data = np.float32(image.reshape(-1, 1))
    gmm.fit(image_data)
    cluster_idx = gmm.predict(image_data)
    clusterCenters = gmm.means_
    max_idx_vec = np.argsort(clusterCenters.flatten())[-selected_components:]
    idx_keep = np.isin(cluster_idx, max_idx_vec)
    cluster_idx[~idx_keep] = 0
    out_put_image = cluster_idx.reshape(image.shape)

    out_put_image = cv2.filter2D(out_put_image, -1, np.ones((12, 12)) / 144)
    out_put_image = maxPooling(out_put_image, [5, 5])

    contours, _ = cv2.findContours(np.uint8(out_put_image), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    boundingBoxes = np.zeros((len(contours), 5))

    for i, contour in enumerate(contours):
        x, y, w, h = cv2.boundingRect(contour)
        boundingBoxes[i, :] = [0, x, y, w, h]

    # 存储bounding box
    folder_path = './results_label/'
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    filePath = f'{folder_path}{os.path.splitext(os.path.basename(pic_path))[0]}.txt'
    np.savetxt(filePath, boundingBoxes, fmt='%d', delimiter=' ')

    Ngt = len(label_matrix)
    Ntt = 0
    Nfa = 0

    for i in range(len(contours)):
        count = 0
        for j in range(len(label_matrix)):
            areaOverlap = xyOverlap(label_matrix[j], boundingBoxes[i, 1:])
            if areaOverlap == 1:
                count += 1
                break
        if count == 0:
            Nfa += 1
        else:
            Ntt += 1

    FoM = Ntt / (Ngt + Nfa)
    Precision = Ntt / (Ntt + Nfa)
    Far = Nfa / (Ntt + Nfa)

    # 显示图像结果并存储
    fig1, axes1 = plt.subplots(2, 1)
    axes1[0].imshow(image_init, cmap='gray')
    axes1[0].set_title('原始图像')
    axes1[1].imshow(image_init, cmap='gray')
    for i in range(len(contours)):
        axes1[1].add_patch(Rectangle((boundingBoxes[i, 1], boundingBoxes[i, 2]), boundingBoxes[i, 3], boundingBoxes[i, 4], edgecolor='r', linewidth=2, fill=False))
    axes1[1].set_title('检测结果')
    axes1[0].axis('off')
    axes1[1].axis('off')
    plt.axis('off')
    plt.tight_layout()

    folder_path = './results_pic/'
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    plt.savefig(f'{folder_path}{os.path.splitext(os.path.basename(pic_path))[0]}.jpg', bbox_inches='tight', pad_inches= 0.0)
    plt.close(fig1)

    fig2 = plt.figure()
    plt.imshow(image_init, cmap='gray')
    for i in range(len(contours)):
        plt.gca().add_patch(Rectangle((boundingBoxes[i, 1], boundingBoxes[i, 2]), boundingBoxes[i, 3], boundingBoxes[i, 4], edgecolor='r', linewidth=2, fill=False))
    plt.axis('off')
    folder_path = './results_pic_original/'
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    plt.savefig(f'{folder_path}{os.path.splitext(os.path.basename(pic_path))[0]}.jpg', bbox_inches='tight', pad_inches = 0.0)
    absolute_path = os.path.abspath(f'{folder_path}{os.path.splitext(os.path.basename(pic_path))[0]}.jpg')
    print(absolute_path)

if __name__ == "__main__":

    if len(sys.argv) != 3:
        sys.exit(1)

    pic_path = sys.argv[1]
    label_path = sys.argv[2]
    ship_detection_fun(pic_path, label_path)