feature_extractor.py

import torch
import numpy as np
import torchvision.models as models
from torchvision import transforms
from PIL import Image
import cv2
from skimage.feature import local_binary_pattern

class FeatureExtractor:
    def __init__(self, device):
        self.device=device
        model = models.mobilenet_v2(weights='IMAGENET1K_V1')
        self.extractor = torch.nn.Sequential(
            model.features,
            torch.nn.AdaptiveAvgPool2d((1, 1)),
            torch.nn.Flatten()
        )
        self.extractor.eval()
        self.extractor = self.extractor.to(device)

        self.preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225]
            )
        ])

    def extract_features(self, img_paths, labels):
        all_features=[]
        new_labels=[]
        total = len(img_paths)
        for i, (path,label) in enumerate(zip(img_paths,labels)):
            try:
                img = Image.open(str(path)).convert("RGB")
            except:
                continue
            img_tensor=self.preprocess(img).unsqueeze(0).to(self.device)
            with torch.no_grad():
                cnn_features = self.extractor(img_tensor).squeeze().cpu().numpy()
            img_gray = np.array(img.convert("L"))
            lbp_feat = extract_lbp(img_gray)
            edge_feat = extract_edge_features(img_gray)
            color_feat =extract_color_stats(img)

            features = np.concatenate([cnn_features, lbp_feat, edge_feat, color_feat])
            all_features.append(features)
            new_labels.append(label)

            if (i + 1) % 100 == 0:
                print(f"  proessed extracting features: {i + 1}/{total}")
        return all_features, new_labels
    
def extract_lbp(img_gray, P=24, R=3):
    lbp = local_binary_pattern(img_gray, P, R, method="uniform")
    hist, _ = np.histogram(lbp, bins=np.arange(0, P + 3), density=True)
    return hist

def extract_edge_features(img_gray):
    edges = cv2.Canny(img_gray, 100, 200)
    edge_ratio = np.sum(edges > 0) / edges.size
    mean_edge = edges.mean() / 255.0
    return np.array([edge_ratio, mean_edge])

def extract_color_stats(img_rgb):
    img_np = np.array(img_rgb) / 255.0
    means = img_np.mean(axis=(0, 1))
    stds = img_np.std(axis=(0, 1))
    return np.concatenate([means, stds])