app.py

import numpy as np
import cv2
import os
import gradio as gr
from skimage.feature import graycomatrix, graycoprops, local_binary_pattern
from skimage.io import imread
from skimage.color import rgb2gray
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.preprocessing import StandardScaler

# Paths to dataset
DATASET_PATH = ""
CATEGORIES = ["stone", "brick", "wood"]

def preprocess_image(image):
    """Ensure the image is 2D grayscale and integer type."""
    if image is None:
        print("❌ Warning: Received an empty image.")
        return None

    # Remove extra batch dimensions if they exist
    if len(image.shape) == 4:
        print(f"⚠️ Removing extra dimension: {image.shape}")
        image = np.squeeze(image, axis=0)

    # Convert RGBA to RGB if necessary
    if len(image.shape) == 3 and image.shape[2] == 4:
        print(f"⚠️ Converting RGBA to RGB: {image.shape}")
        image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)

    # Convert to grayscale if needed
    if len(image.shape) == 3 and image.shape[2] == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Convert to uint8
    return image.astype(np.uint8)

def extract_glcm_features(image):
    """Extract GLCM texture features from a grayscale image."""
    gray = preprocess_image(image)
    if gray is None:
        return None

    glcm = graycomatrix(gray, distances=[1, 3, 5], angles=[0, np.pi/4, np.pi/2, 3*np.pi/4],
                        levels=256, symmetric=True, normed=True)
    contrast = graycoprops(glcm, 'contrast').flatten()
    correlation = graycoprops(glcm, 'correlation').flatten()
    energy = graycoprops(glcm, 'energy').flatten()
    homogeneity = graycoprops(glcm, 'homogeneity').flatten()

    return np.hstack([contrast, correlation, energy, homogeneity])

def extract_lbp_features(image):
    """Extract LBP texture features from a grayscale image."""
    gray = preprocess_image(image)
    if gray is None:
        return None

    radius = 1
    points = 8 * radius
    lbp = local_binary_pattern(gray, P=points, R=radius, method="uniform")
    
    # Compute histogram
    n_bins = int(lbp.max() + 1)
    hist, _ = np.histogram(lbp.ravel(), bins=n_bins, range=(0, n_bins), density=True)

    return hist

def augment_image(image):
    """Perform multiple augmentation techniques on an image."""
    augmented = []

    # Flip horizontally
    augmented.append(cv2.flip(image, 1))

    # Flip vertically
    augmented.append(cv2.flip(image, 0))

    # Rotate 90 degrees clockwise
    augmented.append(cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE))

    # Rotate 90 degrees counterclockwise
    augmented.append(cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE))

    # Add Gaussian noise
    noise = np.random.normal(0, 10, image.shape).astype(np.uint8)
    noisy_image = cv2.add(image, noise)
    augmented.append(noisy_image)

    return augmented

def load_dataset():
    """Load images and extract features while handling different image formats."""
    features_glcm, labels_glcm = [], []
    features_lbp, labels_lbp = [], []
    original_images = []  # Store original images for augmentation

    for label, category in enumerate(CATEGORIES):
        category_path = os.path.join(DATASET_PATH, category)

        def is_image_file(file):
            return file.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".tiff"))

        image_files = [f for f in os.listdir(category_path) if is_image_file(f)]

        for img_file in image_files:
            img_path = os.path.join(category_path, img_file)
            image = imread(img_path)

            if image is None:
                print(f"❌ Skipping: {img_path} (Failed to load)")
                continue

            image = preprocess_image(image)
            original_images.append(image)  # Store for augmentation

            glcm_features = extract_glcm_features(image)
            lbp_features = extract_lbp_features(image)

            if glcm_features is not None and lbp_features is not None:
                features_glcm.append(glcm_features)
                labels_glcm.append(label)

                features_lbp.append(lbp_features)
                labels_lbp.append(label)

    max_len_glcm = max(len(x) for x in features_glcm)
    features_glcm = np.array([np.pad(x, (0, max_len_glcm - len(x)), mode='constant') for x in features_glcm])

    max_len_lbp = max(len(x) for x in features_lbp)
    features_lbp = np.array([np.pad(x, (0, max_len_lbp - len(x)), mode='constant') for x in features_lbp])

    return np.array(features_glcm), np.array(labels_glcm), np.array(features_lbp), np.array(labels_lbp), original_images

# Load dataset
X_glcm, y_glcm, X_lbp, y_lbp, original_images = load_dataset()

scaler_lbp = StandardScaler()
X_lbp = scaler_lbp.fit_transform(X_lbp)

# Create train-test split while keeping track of original indices
X_train_glcm, X_test_glcm, y_train_glcm, y_test_glcm, train_indices, test_indices = train_test_split(
    X_glcm, y_glcm, np.arange(len(y_glcm)), test_size=0.3, random_state=42
)

X_train_lbp, X_test_lbp, y_train_lbp, y_test_lbp = train_test_split(
    X_lbp, y_lbp, test_size=0.3, random_state=42
)

# Apply augmentation only to training images
X_train_glcm_aug, y_train_glcm_aug = list(X_train_glcm), list(y_train_glcm)
X_train_lbp_aug, y_train_lbp_aug = list(X_train_lbp), list(y_train_lbp)

# 🔹 Corrected: Use indices from the original dataset correctly
for idx, original_idx in enumerate(train_indices):
    image = original_images[original_idx]  # Retrieve correct original image
    augmented_images = augment_image(image)

    for aug_img in augmented_images:
        glcm_features = extract_glcm_features(aug_img)
        lbp_features = extract_lbp_features(aug_img)

        if glcm_features is not None and lbp_features is not None:
            X_train_glcm_aug.append(glcm_features)
            y_train_glcm_aug.append(y_train_glcm[idx])  # ✅ Use correct training label

            X_train_lbp_aug.append(lbp_features)
            y_train_lbp_aug.append(y_train_lbp[idx])  # ✅ Use correct training label

# Convert augmented dataset back to numpy arrays
X_train_glcm = np.array(X_train_glcm_aug)
y_train_glcm = np.array(y_train_glcm_aug)
X_train_lbp = np.array(X_train_lbp_aug)
y_train_lbp = np.array(y_train_lbp_aug)


# Train classifiers
svm_glcm = SVC(kernel='linear')
svm_glcm.fit(X_train_glcm, y_train_glcm)
knn_glcm = KNeighborsClassifier(n_neighbors=3)
knn_glcm.fit(X_train_glcm, y_train_glcm)

svm_lbp = SVC(kernel='linear')
svm_lbp.fit(X_train_lbp, y_train_lbp)
knn_lbp = KNeighborsClassifier(n_neighbors=5, metric='manhattan')
knn_lbp.fit(X_train_lbp, y_train_lbp)


# Evaluate models
print("SVM GLCM Performance:\n", classification_report(y_test_glcm, svm_glcm.predict(X_test_glcm)))
print("KNN GLCM Performance:\n", classification_report(y_test_glcm, knn_glcm.predict(X_test_glcm)))
print("SVM LBP Performance:\n", classification_report(y_test_lbp, svm_lbp.predict(X_test_lbp)))
print("KNN LBP Performance:\n", classification_report(y_test_lbp, knn_lbp.predict(X_test_lbp)))


def classify_texture(image, algorithm):
    image = preprocess_image(image)
    
    if algorithm == "SVM (GLCM)":
        features = extract_glcm_features(image).reshape(1, -1)
        prediction = svm_glcm.predict(features)
    elif algorithm == "KNN (GLCM)":
        features = extract_glcm_features(image).reshape(1, -1)
        prediction = knn_glcm.predict(features)
    elif algorithm == "SVM (LBP)":
        features = extract_lbp_features(image).reshape(1, -1)
        features = scaler_lbp.transform(features)
        prediction = svm_lbp.predict(features)
    elif algorithm == "KNN (LBP)":
        features = extract_lbp_features(image).reshape(1, -1)
        features = scaler_lbp.transform(features)
        prediction = knn_lbp.predict(features)
    
    return CATEGORIES[prediction[0]]


interface = gr.Interface(
    fn=classify_texture,
    inputs=[
        gr.Image(type="numpy"),
        gr.Radio(["SVM (GLCM)", "KNN (GLCM)", "SVM (LBP)", "KNN (LBP)"], label="Select Algorithm")
    ],
    outputs=gr.Label(label="Predicted Texture"),
    title="Texture Classification",
    description="Upload an image of a texture and choose an algorithm to classify it as stone, brick, or wood."
)

interface.launch()