Welding-Parts-Anomaly-Detection/small_image_match.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/6/17 0017 上午 10:19
# @Author  : liudan
# @File    : small_image_match.py
# @Software: pycharm

import cv2 as cv
import cv2
import numpy as np
from loguru import logger
import os

from superpoint_superglue_deployment import Matcher
from datetime import datetime
import random
# timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')


def registration(ref_image_path, query_image_path):
    ref_image = cv2.imread(ref_image_path)
    query_image = cv2.imread(query_image_path)

    ref_gray = cv2.cvtColor(ref_image, cv2.COLOR_BGR2GRAY)
    query_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)
    query_gray_rotate = cv2.rotate(query_gray, cv2.ROTATE_180)


    superglue_matcher = Matcher(
        {
            "superpoint": {
                "input_shape": (-1, -1),
                "keypoint_threshold": 0.003,
            },
            "superglue": {
                "match_threshold": 0.5,
            },
            "use_gpu": True,
        }
    )
    _, _, _, _, matches1 = superglue_matcher.match(query_gray, ref_gray)
    _, _, _, _, matches2 = superglue_matcher.match(query_gray_rotate, ref_gray)

    if len(matches1) > len(matches2):
        flag = True
    else:
        flag = False
    print(flag)


    # num_matches = len(matches2)
    # score = 0
    # if num_matches < 4:
    #     print('The welding has been done in reverse')
    #     score = 0
    # else:
    #     score = 1
    # print("score:", score)

if __name__ == "__main__":
    ref_image_path = './data/small_image_match/json_image1_E_0.jpg'
    query_image_path = './data/small_image_match/json_image2_E_0.jpg'
    registration(ref_image_path, query_image_path)