Upload 3 files

predict_database.py

@@ -0,0 +1,13 @@
+import os
+
+from segment_key import final_features
+
+image_dir = '../augmentation/testing/3'
+
+for filename in os.listdir(image_dir):
+    if filename.endswith('.jpg'):
+        image_path = os.path.join(image_dir, filename)
+        features = final_features(image_path)
+        with open('./prediction/database.txt', 'a+') as f:
+            f.write(filename + ';' + str(features) + '\n')
+            print('successfully predicted ' + filename)

segment_key.py

@@ -0,0 +1,344 @@
+import math
+
+import cv2
+import imutils
+import numpy as np
+from sklearn.preprocessing import MinMaxScaler
+from ultralytics import YOLO
+
+from models.birefnet import BiRefNet
+from util.utils import check_state_dict
+from PIL import Image
+import torch
+from torchvision import transforms
+from openvino.runtime import Core
+
+model = BiRefNet(bb_pretrained=False)
+state_dict = torch.load('models/weights/BiRefNet-general-epoch_244.pth', map_location='cpu')
+state_dict = check_state_dict(state_dict)
+model.load_state_dict(state_dict)
+
+# Input Data
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+import torch
+from PIL import Image
+import torchvision.transforms as transforms
+
+
+def export_to_onnx(model, dummy_input, onnx_path='model.onnx'):
+    # Export the PyTorch model to ONNX format
+    torch.onnx.export(model, dummy_input, onnx_path, verbose=True, input_names=['input'], output_names=['output'])
+
+
+image = Image.open('./examples/img.jpg')
+dummy_input = torch.randn(1, 3, 224, 224)  # Adjust input size to match your model
+export_to_onnx(model, image)
+
+
+def pred_segmentation(image, box=[-1, -1, -1, -1]):
+    core = Core()
+    model_path = 'model_ir/model.xml'
+    compiled_model = core.compile_model(model_path, device_name='GPU')  # Use "GPU" to target Intel GPU
+    print('predicting segmentation...')
+
+    # box: left, top, right, bottom
+    w, h = image.size[:2]
+
+    # Adjust box coordinates if necessary
+    for idx_coord_value, coord_value in enumerate(box):
+        if coord_value == -1:
+            box[idx_coord_value] = [0, 0, w, h][idx_coord_value]
+
+    # Crop the image based on the box
+    image_crop = image.crop(box)
+
+    # Transform the image
+    input_image = transform_image(image_crop)
+    input_image = np.expand_dims(input_image, axis=0)  # Add batch dimension
+
+    # Run inference
+    infer_request = compiled_model.create_infer_request()
+    result = infer_request.infer(inputs={'input': input_image})
+
+    pred = result['output'][0]  # Adjust if output name is different
+
+    # Post-process predictions
+    canvas = np.zeros_like(pred)
+
+    # Calculate the bounding box to place the result back on the canvas
+    box_to_canvas = [
+        int(round(coord_value * (canvas.shape[-1] / w, canvas.shape[-2] / h)[idx_coord_value % 2]))
+        for idx_coord_value, coord_value in enumerate(box)
+    ]
+
+    pred = np.resize(pred, (box_to_canvas[3] - box_to_canvas[1], box_to_canvas[2] - box_to_canvas[0]))
+
+    canvas[box_to_canvas[1]:box_to_canvas[3], box_to_canvas[0]:box_to_canvas[2]] = pred
+
+    # Convert the canvas to PIL image for visualization
+    pred_pil = Image.fromarray((canvas * 255).astype(np.uint8))  # Rescale to [0, 255] for visualization
+
+    return pred_pil
+
+
+def pred_bbox(image_path):
+    print('predicting bounding box...')
+    image = cv2.imread(image_path)
+    model = YOLO('models/weights/yolo_finetuned.pt')
+
+    # Perform prediction
+    results = model(image)
+    boxes = results[0].boxes.xyxy.cpu().numpy()[0]
+
+    # Extract the bounding box coordinates
+    x1, y1, x2, y2 = map(int, list(boxes))
+    return [x1, y1, x2, y2]
+
+
+def get_kps_from_pil(pil_image):
+    print('converting keypoints...')
+    image_array = np.array(pil_image)
+
+    # Find contours using OpenCV
+    contours, _ = cv2.findContours(image_array, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    # Find the largest contour by area
+    largest_contour = max(contours, key=cv2.contourArea)
+    largest_contour = np.array(largest_contour)
+    contour = []
+    for i in range(len(largest_contour)):
+        contour.append(largest_contour[i][0])
+    scaler = MinMaxScaler()
+    kps = scaler.fit_transform(contour)
+    kps = np.array(kps)
+    kps = kps * 299
+    kps = np.int32(kps)
+    return kps
+
+
+def get_features_up(contour):
+    feature = []
+    for i in range(0, 300):
+        position = 0
+        unsorted_features = []
+        for j in range(len(contour)):
+            point = contour[j]
+            prev_point = point
+            if j != 0:
+                prev_point = contour[j - 1]
+            if point[0] > i and position == 0:
+                position = 1
+            elif point[0] < i and position == 0:
+                position = -1
+            elif point[0] > i and position == -1:
+                unsorted_features.append((point[1] + prev_point[1]) // 2)
+                position = 1
+            elif point[0] < i and position == 1:
+                position = -1
+                unsorted_features.append((point[1] + prev_point[1]) // 2)
+            elif point[0] == i and position == 1:
+                unsorted_features.append(point[1])
+                position = -1
+            elif point[0] == i and position == -1:
+                unsorted_features.append(point[1])
+                position = 1
+            elif point[0] == i and position == 0:
+                position = 1
+
+        if len(unsorted_features) != 0:
+            if len(unsorted_features) == 1:
+                unsorted_features.append((contour[0][1] + contour[-1][1]) // 2)
+            unsorted_features.sort()
+            feature.append(max(unsorted_features))
+        else:
+            feature.append(-1)
+
+    return feature
+
+
+def get_features_down(contour):
+    feature = []
+    for i in range(0, 300):
+        position = 0
+        unsorted_features = []
+        for j in range(len(contour)):
+            point = contour[j]
+            prev_point = point
+            if j != 0:
+                prev_point = contour[j - 1]
+            if point[0] > i and position == 0:
+                position = 1
+            elif point[0] < i and position == 0:
+                position = -1
+            elif point[0] > i and position == -1:
+                unsorted_features.append((point[1] + prev_point[1]) // 2)
+                position = 1
+            elif point[0] < i and position == 1:
+                position = -1
+                unsorted_features.append((point[1] + prev_point[1]) // 2)
+            elif point[0] == i and position == 1:
+                unsorted_features.append(point[1])
+                position = -1
+            elif point[0] == i and position == -1:
+                unsorted_features.append(point[1])
+                position = 1
+            elif point[0] == i and position == 0:
+                position = 1
+
+        if len(unsorted_features) != 0:
+            if len(unsorted_features) == 1:
+                unsorted_features.append((contour[0][1] + contour[-1][1]) // 2)
+            unsorted_features.sort()
+            feature.append(min(unsorted_features))
+        else:
+            feature.append(-1)
+
+    return feature
+
+
+def get_features_right(contour):
+    feature = []
+    for i in range(0, 300):
+        position = 0
+        unsorted_features = []
+        for j in range(len(contour)):
+            point = contour[j]
+            prev_point = point
+            if j != 0:
+                prev_point = contour[j - 1]
+            if point[1] > i and position == 0:
+                position = 1
+            elif point[1] < i and position == 0:
+                position = -1
+            elif point[1] > i and position == -1:
+                unsorted_features.append((point[0] + prev_point[0]) // 2)
+                position = 1
+            elif point[1] < i and position == 1:
+                position = -1
+                unsorted_features.append((point[0] + prev_point[0]) // 2)
+            elif point[1] == i and position == 1:
+                unsorted_features.append(point[0])
+                position = -1
+            elif point[1] == i and position == -1:
+                unsorted_features.append(point[0])
+                position = 1
+            elif point[1] == i and position == 0:
+                position = 1
+
+        if len(unsorted_features) != 0:
+            if len(unsorted_features) == 1:
+                unsorted_features.append((contour[0][0] + contour[-1][0]) // 2)
+            unsorted_features.sort()
+            feature.append(min(unsorted_features))
+        else:
+            feature.append(-1)
+
+    return feature
+
+
+def get_features_left(contour):
+    feature = []
+    for i in range(0, 300):
+        position = 0
+        unsorted_features = []
+        for j in range(len(contour)):
+            point = contour[j]
+            prev_point = point
+            if j != 0:
+                prev_point = contour[j - 1]
+            if point[1] > i and position == 0:
+                position = 1
+            elif point[1] < i and position == 0:
+                position = -1
+            elif point[1] > i and position == -1:
+                unsorted_features.append((point[0] + prev_point[0]) // 2)
+                position = 1
+            elif point[1] < i and position == 1:
+                position = -1
+                unsorted_features.append((point[0] + prev_point[0]) // 2)
+            elif point[1] == i and position == 1:
+                unsorted_features.append(point[0])
+                position = -1
+            elif point[1] == i and position == -1:
+                unsorted_features.append(point[0])
+                position = 1
+            elif point[1] == i and position == 0:
+                position = 1
+
+        if len(unsorted_features) != 0:
+            if len(unsorted_features) == 1:
+                unsorted_features.append((contour[0][0] + contour[-1][0]) // 2)
+            unsorted_features.sort()
+            feature.append(max(unsorted_features))
+        else:
+            feature.append(-1)
+
+    return feature
+
+
+def extract_features(contour):
+    print('extracting features...')
+    return get_features_down(contour) + get_features_up(contour) + get_features_right(contour) + get_features_left(contour)
+
+
+def final_features(image_path):
+    image = Image.open(image_path)
+    image = rotate_image(image)
+    pil_image = pred_segmentation(image, pred_bbox(image_path))
+    contour = get_kps_from_pil(pil_image)
+    return extract_features(contour)
+
+
+def predict_kps(image):
+    model = YOLO('models/weights/yolo_finetuned.pt')
+    # Perform prediction
+    results = model(image)
+    kps = results[0].masks.xy[0]
+    return kps
+
+
+def calculate_angle(p1, p2):
+    delta_y = p2[1] - p1[1]
+    delta_x = p2[0] - p1[0]
+    return math.degrees(np.arctan2(delta_y, delta_x))
+
+
+# Function to rotate points by a given angle
+def calculate_square(img):
+    np_image = np.array(img)
+    # Convert RGB (PIL) to BGR (OpenCV)
+    if np_image.ndim == 3:  # Check if the image is colored
+        cv_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
+    else:
+        # For grayscale images, no conversion is needed
+        cv_image = np_image
+
+    rect = cv2.minAreaRect(predict_kps(cv_image))
+    box = cv2.boxPoints(rect)
+    box = np.int32(box)
+    return box
+
+
+def rotate_image(image):
+    square = calculate_square(image)
+    # Calculate the lengths of the sides
+    side_lengths = [np.linalg.norm(square[i] - square[i + 1]) for i in range(len(square) - 1)]
+
+    # Find the indices of the larger side
+    max_index = np.argmax(side_lengths)
+
+    # Find the two points that form the largest side
+    p1, p2 = square[max_index], square[max_index + 1]
+
+    # Calculate the angle between this side and the horizontal axis
+    angle = calculate_angle(p1, p2)
+
+    # Rotate the square to align the largest side with the horizontal axis
+    rotated_image = image.rotate(angle)
+
+    return rotated_image

test.py

@@ -0,0 +1,82 @@
+import os
+from ast import literal_eval
+
+from segment_key import *
+from matplotlib import pyplot as plt
+
+
+def show_kps(contour):
+    list1 = range(0, 310)
+    list2 = list(zip(get_features_right(contour), list1))
+
+    x_coords = [point[0] for point in list2]
+    y_coords = [point[1] for point in list2]
+    plt.scatter(x_coords, y_coords, c='red', marker='o', label='Keypoints')
+
+    list2 = list(zip(get_features_left(contour), list1))
+    x_coords = [point[0] for point in list2]
+    y_coords = [point[1] for point in list2]
+    plt.scatter(x_coords, y_coords, c='red', marker='o', label='Keypoints')
+
+    list2 = list(zip(list1, get_features_up(contour)))
+    x_coords = [point[0] for point in list2]
+    y_coords = [point[1] for point in list2]
+    plt.scatter(x_coords, y_coords, c='red', marker='o', label='Keypoints')
+
+    list2 = list(zip(list1, get_features_down(contour)))
+    x_coords = [point[0] for point in list2]
+    y_coords = [point[1] for point in list2]
+    plt.scatter(x_coords, y_coords, c='red', marker='o', label='Keypoints')
+
+    plt.show()
+
+
+def get_all_features():
+    contours = []
+    with open('prediction/database.txt', 'r') as file:
+        lines = file.readlines()
+        for line in lines:
+            results = (line.split(';')[1])
+            results = results.replace(",,", ",'',")
+            results = literal_eval(results)
+            results = np.array(results)
+            contours.append((line.split(';')[0], results))
+
+    return contours
+
+
+def cos_similarity(feature1, feature2):
+    return np.dot(feature1, feature2) / (np.linalg.norm(feature1) * np.linalg.norm(feature2))
+
+
+def predict_match(image_path):
+    main_name = os.path.basename(image_path)[:-11] + '.jpg'
+    main_feature = final_features(image_path)
+    contours = get_all_features()
+
+    l = []
+
+    for image in contours:
+        feature = image[1]
+        feature_similarity = 1 - cos_similarity(feature, main_feature)
+
+        l.append([image[0], feature_similarity])
+
+    l.sort(key=lambda x: x[1])
+
+    print(l)
+    print(l[0])
+    index_in_list = -1
+    for i in range(len(l)):
+        if l[i][0] == main_name:
+            index_in_list = i
+    return l[0][0], l[0][1], index_in_list
+
+
+# for image_file in os.listdir('../augmentation/testing/3'):
+#     image_path = os.path.join('../augmentation/testing/3', image_file)
+#     best_image_name, best_match_value, index = predict_match(image_path)
+#     with open('performance/performance_results.txt', 'a+') as file:
+#         file.write(image_file + ';' + best_image_name + ';' + str(best_match_value) + ';' + str(index) + '\n')
+
+predict_match('./examples/img.jpg')