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):
    """
    Extract features from the image and detection result.
    
    Args:
        image: PIL Image
        detection_result: Dictionary containing detection results
        
    Returns:
        Dictionary of features
    """
    np_image = np.array(image)
    amplitude = calculate_amplitude(np_image)
    amplitude_class = classify_amplitude(amplitude)

    # Handle box calculation
    if detection_result and 'boxes' in detection_result and len(detection_result['boxes']) > 0:
        box = detection_result['boxes'][0]
        if not isinstance(box, list):
            box = box.tolist()
        distribution_range = calculate_distribution_range(box)
        distribution_class = classify_distribution_range(distribution_range)
    else:
        distribution_range = 0
        distribution_class = "小"

    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,
        "振幅值": float(amplitude),
        "分布范围值": float(distribution_range),
        "衰减速度值": float(attenuation_rate),
        "反射次数值": int(reflection_count)
    }