Welding-Parts-Anomaly-Detection/demo.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/5/23 0023 下午 2:25
# @Author  : liudan
# @File    : demo.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')
#import image_similarity_count


def registration_demo():
    image_dir = './data/yongsheng_image/3_image_query'
    #images = [f for f in os.listdir(image_dir) if os.path.isfile(os.path.join(image_dir, f))]
    demo_image_path = './data/yongsheng_image/3_regis_result/'
    os.makedirs(demo_image_path, exist_ok=True)

    wraped_image_list =[]
    for filename in os.listdir(image_dir):
        # 检查文件是否是图片（这里假设是.jpg格式）
        if filename.endswith('.JPG') or filename.endswith('.jpg'):
            # 构建图片的完整路径
            image1_path = os.path.join(image_dir, filename)
            query_image = cv2.imread(image1_path)
            #query_image = cv2.imread('./data/images/frame_002100_small.jpg') #原图2160*3840
            ref_image = cv2.imread('./data/yongsheng_image/ref_image/DSC_0451.JPG')

            # new_size = (1420, 800) # W,H
            # query_image_resize = cv2.resize(query_image,new_size)
            # ref_image_resize = cv2.resize(ref_image,new_size)
            height, width = ref_image.shape[:2]
            scale_factor = 0.37
            query_image_resize = cv2.resize(query_image, dsize=None, fx=scale_factor, fy=scale_factor)
            ref_image_resize = cv2.resize(ref_image, dsize=None, fx=scale_factor, fy=scale_factor)

            query_gray = cv2.cvtColor(query_image_resize, cv2.COLOR_BGR2GRAY)
            ref_gray = cv2.cvtColor(ref_image_resize, cv2.COLOR_BGR2GRAY)

            superglue_matcher = Matcher(
                {
                    "superpoint": {
                        "input_shape": (-1, -1),
                        "keypoint_threshold": 0.003,
                    },
                    "superglue": {
                        "match_threshold": 0.5,
                    },
                    "use_gpu": True,
                }
            )
            query_kpts, ref_kpts, _, _, matches = superglue_matcher.match(query_gray, ref_gray)
            M, mask = cv2.findHomography(
                np.float64([query_kpts[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2),
                np.float64([ref_kpts[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2),
                method=cv2.USAC_MAGSAC,
                ransacReprojThreshold=5.0,
                maxIters=10000,
                confidence=0.95,
            )
            logger.info(f"number of inliers: {mask.sum()}")
            matches = np.array(matches)[np.all(mask > 0, axis=1)]
            matches = sorted(matches, key=lambda match: match.distance)
            matched_image = cv2.drawMatches(
                query_image_resize,
                query_kpts,
                ref_image_resize,
                ref_kpts,
                matches[:50],
                None,
                flags=2,
            )

            #demo_image_path = './data/yongsheng_image/1_regis_result/'
            match_file_name = f"match_image_{filename}.jpg"
            cv2.imwrite(os.path.join(demo_image_path , match_file_name), matched_image)
            # cv2.imwrite(os.path.join(demo_image_path + f"matched_image.jpg"), matched_image)

            wrap_image = cv.warpPerspective(query_image_resize, M, (ref_image_resize.shape[1], ref_image_resize.shape[0]))
            wrap_image = cv2.resize(wrap_image,(ref_image.shape[1], ref_image.shape[0]))
            wrap_filename = f"wrap_image_{filename}.jpg"
            cv2.imwrite(os.path.join(demo_image_path, wrap_filename), wrap_image)
            #cv2.imwrite(os.path.join(demo_image_path + f"wrap_image.jpg"), wrap_image)
            wraped_image_list.append((wrap_image,wrap_filename))

            result_image = cv2.subtract(ref_image, wrap_image)
            result_file_name = f"result_image_{filename}.jpg"
            #cv2.imwrite(os.path.join(demo_image_path + f"result_image.jpg"), result_image)
            cv2.imwrite(os.path.join(demo_image_path, result_file_name), result_image)

    return wraped_image_list



if __name__ == "__main__":
    registration_demo()

    #image_similarity_count.compare_boxes_similarity(image1_path, image2_path, json_file_path, similarity_threshold=0.5)