preprocess_quality.py

import cv2
import numpy as np
from PIL import Image

class GeometricFeatureExtractor:
    def __init__(self, image_size=128):
        self.image_size = image_size
        self.last_feature_dict = None  

    def extract_features(self, pil_image: Image.Image):
        img = np.array(pil_image.convert("L"))
        img = cv2.resize(img, (self.image_size, self.image_size))
        blur = cv2.GaussianBlur(img, (5,5),0)
        _, thresh = cv2.threshold(
            blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
        )

        contours, _ = cv2.findContours(
            thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        if not contours:
            return None
        
        contour = max(contours, key=cv2.contourArea)
        area = cv2.contourArea(contour)
        perimeter = cv2.arcLength(contour, True)
        if area == 0 or perimeter == 0:
            return None

        x, y, w, h = cv2.boundingRect(contour)
        aspect_ratio = w / h if h != 0 else 0

        if len(contour) >= 5:
            ellipse = cv2.fitEllipse(contour)
            major_axis = max(ellipse[1])
            minor_axis = min(ellipse[1])
            eccentricity = (
                np.sqrt(1 - (minor_axis / major_axis)**2)
                if major_axis > 0 else 0
            )
        else:
            eccentricity = 0

        compactness = (perimeter**2) / (4*np.pi*area)
        circularity = (4*np.pi*area) / (perimeter**2)

        epsilon = 0.02 * perimeter
        approx = cv2.approxPolyDP(contour, epsilon, True)
        num_vertices = len(approx)

        angle_error = 0.0
        if num_vertices >= 3:
            angles = []
            for i in range(num_vertices):
                p1 = approx[i % num_vertices][0]
                p2 = approx[(i+1) % num_vertices][0]
                p3 = approx[(i+2) % num_vertices][0]

                v1 = p1 - p2
                v2 = p3 - p2

                cosang = np.dot(v1, v2) / (
                    np.linalg.norm(v1)*np.linalg.norm(v2) + 1e-6
                )
                angle = np.degrees(np.arccos(np.clip(cosang, -1, 1)))
                angles.append(angle)

            ideal = 60 if num_vertices == 3 else 90
            angle_error = float(np.mean(np.abs(np.array(angles) - ideal)))

        side_lengths = []
        for i in range(num_vertices):
            p1 = approx[i][0]
            p2 = approx[(i+1) % num_vertices][0]
            side_lengths.append(np.linalg.norm(p1 - p2))

        side_length_variance = (
            np.std(side_lengths) / (np.mean(side_lengths) + 1e-6)
            if len(side_lengths) >= 2 else 0
        )

        moments = cv2.moments(contour)
        hu_moments = cv2.HuMoments(moments).flatten()

        edges = cv2.Canny(thresh, 50, 150)
        edge_density = np.sum(edges > 0) / (self.image_size**2)

        hull = cv2.convexHull(contour)
        hull_area = cv2.contourArea(hull)
        solidity = area / hull_area if hull_area > 0 else 0

        self.last_feature_dict = {
            "area_norm": area / (self.image_size**2),
            "perimeter_norm": perimeter / (self.image_size*4),
            "aspect_ratio": aspect_ratio,
            "eccentricity": eccentricity,
            "compactness": compactness,
            "circularity": circularity,
            "num_vertices": num_vertices,
            "edge_density": edge_density,
            "solidity": solidity,
            "angle_error": angle_error,
            "side_length_variance": side_length_variance
        }

        features = np.array([
            self.last_feature_dict["area_norm"],
            self.last_feature_dict["perimeter_norm"],
            aspect_ratio,
            eccentricity,
            compactness,
            circularity,
            num_vertices,
            edge_density,
            solidity,
            *hu_moments
        ], dtype=np.float32)

        return features