functions.py

import numpy as np
import matplotlib.pyplot as plt

from typing import *
import cv2

import torch


def extract_sift_features(image: np.ndarray) -> Tuple[Sequence[cv2.KeyPoint], cv2.UMat]:
    """
    Extracts SIFT descriptors from an image.

    Parameters:
    -----------
    image (numpy.ndarray): The input image.

    Returns:
    --------
    Tuple[Sequence[cv2.KeyPoint], cv2.UMat]: A tuple containing the keypoints and descriptors.
    """

    # Create a SIFT detector
    sift = cv2.SIFT.create()

    # Detect keypoints and compute descriptors
    keypoints, descriptors = sift.detectAndCompute(image, None)

    return keypoints, descriptors


def match_descriptors(
    descriptors_src: cv2.UMat, descriptors_tgt: cv2.UMat, max_ratio: float = 0.8
) -> Tuple[Sequence[Sequence[cv2.DMatch]], Sequence[Sequence[cv2.DMatch]]]:
    """
    Matches descriptors from the test image to the training descriptors using brute-force
    matches and Lowe's ratio test.

    Returns the good matches and all matches.

    Parameters:
    -----------
    descriptors_src (cv2.UMat): The source image descriptors.
    descriptors_tgt (cv2.UMat): The target image descriptors.
    max_ratio (float): The maximum ratio for Lowe's ratio test.

    Returns:
    --------
    Tuple[Sequence[Sequence[cv2.DMatch]], Sequence[Sequence[cv2.DMatch]]]: A tuple containing the good matches and all matches.

    """
    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=False)
    matches = bf.knnMatch(descriptors_src, descriptors_tgt, k=2)

    # Apply Lowe's ratio test
    good_matches = []

    # Iterate through the matches and apply the ratio test
    # selecting only the good matches
    for i, (m, n) in enumerate(matches):
        if m.distance < max_ratio * n.distance:
            good_matches.append([m])

    return good_matches, matches


def partial_gradient(w, w0, example_train, label_train):
    """
    Computes the derivatives of the partial loss Li with respect to each of the classifier parameters.

    Parameters:
    - w: weight vector (same shape as example_train)
    - w0: bias term (scalar)
    - example_train: input image / example (same shape as w)
    - label_train: 0 or 1 (negative or positive example)

    Returns:
    - wgrad: gradient with respect to w
    - w0grad: gradient with respect to w0
    """
    y = np.sum(example_train * w) + w0
    p = np.exp(y) / (1 + np.exp(y))

    wgrad = (p - label_train) * example_train
    w0grad = p - label_train

    return wgrad, w0grad


def process_epoch(w, w0, lrate, examples_train, labels_train, random_order=True):
    """
    Performs one epoch of stochastic gradient descent.

    Parameters:
    - w: weight array (same shape as examples)
    - w0: bias term (scalar)
    - lrate: learning rate (scalar)
    - examples_train: list or array of training examples (e.g., shape (N, 35, 35))
    - labels_train: array of labels (shape (N,))

    Returns:
    - Updated w and w0 after one epoch
    """

    if random_order:
        idx = np.random.permutation(len(examples_train))
    else:
        idx = np.arange(len(examples_train))
    examples_train = examples_train[idx]
    labels_train = labels_train[idx]

    for i in range(len(examples_train)):
        [wgrad, w0grad] = partial_gradient(w, w0, examples_train[i], labels_train[i])
        w = w - lrate * wgrad
        w0 = w0 - lrate * w0grad
        # Plot the weight matrix every 100 iterations
        if i % 100 == 0:
            f, ax = plt.subplots()
            ax.imshow(w, cmap="gray")
            ax.set_title(f"Epoch {i+1}")
            plt.show()

    return w, w0


def train_model(
    model,
    train_loader,
    test_loader,
    num_epochs,
    optimizer,
    loss_fn,
    device,
    debug=False,
):
    for epoch in range(num_epochs):
        model.train()  # Set the model to training mode
        running_loss = 0.0
        correct = 0
        total = 0

        for inputs, targets in train_loader:
            # Move data to the correct device (GPU or CPU)
            inputs, targets = inputs.to(device), targets.to(device)

            # Zero the parameter gradients
            optimizer.zero_grad()

            # Forward pass
            outputs = model(inputs)

            # Compute the loss
            loss = loss_fn(outputs, targets)

            # Backward pass and optimize
            loss.backward()
            optimizer.step()

            # Update statistics
            running_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total += targets.size(0)
            correct += (predicted == targets).sum().item()

        # Print the training statistics
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100 * correct / total

        if debug:
            print(
                f"Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%"
            )

        # Validation loop (optional, to check model performance on validation data)
        model.eval()  # Set the model to evaluation mode
        val_correct = 0
        val_total = 0
        with torch.no_grad():  # Disable gradient calculation for validation
            for inputs, targets in test_loader:
                inputs, targets = inputs.to(device), targets.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs, 1)
                val_total += targets.size(0)
                val_correct += (predicted == targets).sum().item()

        val_acc = 100 * val_correct / val_total
        if debug:
            print(f"Validation Accuracy: {val_acc:.2f}%")

    if debug:
        print("Training complete.")
    return model