feature_extraction.py

import numpy as np
from scipy import signal
from config import AMPLITUDE_THRESHOLD, DISTRIBUTION_THRESHOLD, ATTENUATION_THRESHOLD, REFLECTION_THRESHOLD


def calculate_amplitude(image):
    return np.max(image)


def classify_amplitude(amplitude):
    return "强" if amplitude > AMPLITUDE_THRESHOLD else "弱"


def calculate_distribution_range(box):
    width, height = box[2] - box[0], box[3] - box[1]
    return width * height


def classify_distribution_range(area):
    return "大" if area > DISTRIBUTION_THRESHOLD else "小"


def calculate_attenuation_rate(image):
    gradient = np.gradient(np.array(image).mean(axis=2))
    return np.mean(np.abs(gradient))


def classify_attenuation_rate(rate):
    return "快" if rate > ATTENUATION_THRESHOLD else "慢"


def count_reflections(image, prominence=10):
    gray = np.mean(np.array(image), axis=2)
    peaks, _ = signal.find_peaks(np.mean(gray, axis=1), prominence=prominence)
    return len(peaks)


def classify_reflections(count):
    return "多次反射" if count >= REFLECTION_THRESHOLD else "单次反射"


def extract_features(image, detection_result):
    np_image = np.array(image)
    amplitude = calculate_amplitude(np_image)
    amplitude_class = classify_amplitude(amplitude)

    box = detection_result['boxes'][0].tolist()
    distribution_range = calculate_distribution_range(box)
    distribution_class = classify_distribution_range(distribution_range)

    attenuation_rate = calculate_attenuation_rate(np_image)
    attenuation_class = classify_attenuation_rate(attenuation_rate)

    reflection_count = count_reflections(np_image)
    reflection_class = classify_reflections(reflection_count)

    return {
        "振幅": amplitude_class,
        "分布范围": distribution_class,
        "衰减速度": attenuation_class,
        "反射次数": reflection_class,
        "振幅值": amplitude,
        "分布范围值": distribution_range,
        "衰减速度值": attenuation_rate,
        "反射次数值": reflection_count
    }