blockblast_classify.py

import cv2
import numpy as np
import os
import ast


def classify(input_image_path, split_points, split_images_folder):
    # Load the image
    image = cv2.imread(input_image_path)

    # Check if the image was loaded successfully
    if image is None:
        raise FileNotFoundError(
            f"Unable to load image at path: {input_image_path}. Please check the file path and integrity."
        )

    # Resize the image for easier processing
    scale_percent = 50
    new_width = int(image.shape[1] * scale_percent / 100)
    new_height = int(image.shape[0] * scale_percent / 100)
    dim = (new_width, new_height)
    resized_image = cv2.resize(image, dim, interpolation=cv2.INTER_AREA)

    # Crop the region of interest
    y_start = int(1950 * scale_percent / 100)
    y_end = int(2320 * scale_percent / 100)
    cropped_image = resized_image[y_start:y_end, :]

    # Remove pixels from L/R
    if cropped_image.shape[1] > 200:
        cropped_image = cropped_image[:, 50:-50]
    else:
        raise ValueError("Image is too narrow after cropping.")

    # Remove the color rgb(56, 76, 132) from the cropped image with a larger margin
    # Convert RGB to BGR for OpenCV
    lower_bound = np.array([122, 66, 46])
    upper_bound = np.array([142, 86, 66])
    mask = cv2.inRange(cropped_image, lower_bound, upper_bound)
    cropped_image[mask != 0] = [0, 0, 0]

    # Split the cropped image at the specified split points
    height, width = cropped_image.shape[:2]
    x1, x2 = split_points

    # Ensure the split points are within bounds
    if x1 >= width or x2 >= width or x1 <= 0 or x2 <= 0 or x1 >= x2:
        raise ValueError("Invalid split points for the image.")

    split_images = [
        cropped_image[:, :x1],
        cropped_image[:, x1:x2],
        cropped_image[:, x2:],
    ]

    # Iterate over each split image and find the closest match from the split_images_folder
    results = []
    for idx, split_image in enumerate(split_images):
        if split_image.size == 0:
            print(f"Split {idx + 1} of input image is empty.")
            results.append(None)
            continue

        best_match_filename = None
        best_match_score = float("inf")

        # Iterate through all the images in the split_images_folder
        for filename in os.listdir(split_images_folder):
            if filename.endswith((".png", ".jpg", ".jpeg")):
                existing_image_path = os.path.join(split_images_folder, filename)
                existing_image = cv2.imread(existing_image_path, cv2.IMREAD_GRAYSCALE)

                # Check if the existing image was loaded successfully
                if existing_image is None:
                    print(
                        f"Unable to load existing image at path: {existing_image_path}."
                    )
                    continue

                # Resize the current split image to match the existing image
                resized_split_image = cv2.resize(
                    split_image, (existing_image.shape[1], existing_image.shape[0])
                )

                # Convert both images to grayscale
                split_gray = cv2.cvtColor(resized_split_image, cv2.COLOR_BGR2GRAY)

                # Calculate the difference using Mean Squared Error
                diff = cv2.absdiff(split_gray, existing_image)
                mse = np.mean(diff**2)

                # Update the best match if current score is lower
                if mse < best_match_score:
                    best_match_score = mse
                    best_match_filename = filename

        if best_match_filename is not None:
            # Remove any suffix after the list representation (e.g., '_2')
            clean_filename = best_match_filename.split("_")[0].split(".")[0]
            try:
                shape_array = ast.literal_eval(clean_filename)
                results.append(shape_array)
            except (ValueError, SyntaxError):
                print(f"Unable to parse filename as list: {best_match_filename}.")
                results.append(None)
            print(
                f"Split {idx + 1} of input image matches best with: {best_match_filename}, Shape: {shape_array}"
            )
        else:
            results.append(None)
            print(f"No match found for split {idx + 1} of input image.")

    return results


if __name__ == "__main__":
    input_image_path = r"Image.jpg"
    split_images_folder = r"Block Blast Cropped Masked"
    split_points = (175, 400)
    results = classify(input_image_path, split_points, split_images_folder)
    print("Final shape matches:", results)