obstruction_detector.py

import cv2 as cv
import numpy as np
from scipy.signal import find_peaks
from PIL import Image  # Import PIL

class ObstructionDetector:
    def __init__(self, threshold=500):
        self.threshold = threshold

    def preprocess_image(self, image):
        # Convert the image to grayscale if it's a color image
        if len(image.shape) == 3:
            image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)

        # Apply Gaussian blur to reduce noise
        preprocessed_image = cv.GaussianBlur(image, (5, 5), 0)

        # Perform other preprocessing steps as needed (e.g., contrast adjustment, histogram equalization)

        return preprocessed_image

    def plot_histogram(self, image):
        # Calculate the histogram
        histogram = cv.calcHist([image], [0], None, [256], [0, 256])

        # Smoothing the histogram using a simple moving average (window size = 5)
        kernel = np.ones((5, 1)) / 5
        smoothed_histogram = cv.filter2D(histogram, -1, kernel)

        return smoothed_histogram

    def count_histogram_peaks(self, smoothed_histogram):
        # Find peaks in the smoothed histogram with frequency greater than the threshold
        peaks, _ = find_peaks(smoothed_histogram.flatten(), height=self.threshold)
        return peaks

    def detect_obstruction(self, pil_image):  # Accept PIL image directly
        # Convert PIL image to NumPy array
        img = np.array(pil_image)

        # Preprocess the image
        preprocessed_img = self.preprocess_image(img)

        # Count the number of peaks in the smoothed histogram above the threshold
        smoothed_histogram = self.plot_histogram(preprocessed_img)
        peaks = self.count_histogram_peaks(smoothed_histogram)

        # Check if peaks are too close together
        peak_spacing = np.diff(peaks)
        if len(peak_spacing) == 0 or np.all(peak_spacing < 10):
            report = "A imagem NÃO contém obstrução significativa | e NÃO possui múltiplas distribuições de densidade claramente distintas."
        else:
            report = "A imagem contém obstrução significativa | possui múltiplas distribuições de densidade claramente distintas."

        return report