123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143 |
- #!/usr/bin/env python
- # -*- coding: utf-8 -*-
- # @Time : 2024/6/6 0006 上午 9:21
- # @Author : liudan
- # @File : demo_env.py
- # @Software: pycharm
- import cv2 as cv
- import cv2
- import numpy as np
- from loguru import logger
- import os
- import json
- from superpoint_superglue_deployment import Matcher
- from datetime import datetime
- import random
- import yaml
- # timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
- #import image_similarity_count
- def process_image_with_mask(ref_image, scale_factor, json_ref_path, ref_gray):
- with open(json_ref_path, 'r') as f:
- data = json.load(f)
- shapes = data['shapes']
- shape = shapes[0]
- if shape['shape_type'] == 'polygon':
- coords = [(int(round(x * scale_factor)), int(round(y * scale_factor))) for x, y in shape['points']]
- else:
- coords = []
- mask = np.zeros(ref_gray.shape, dtype=np.uint8) * 255
- pts = np.array(coords, np.int32)
- cv2.fillPoly(mask, [pts], 1)
- ref_gray_masked = cv2.bitwise_and(ref_gray, ref_gray, mask=mask)
- # cv2.imwrite('1111111.jpg',ref_gray_masked)
- return ref_gray_masked
- def registration_demo(image_dir,demo_image_path, json_ref_path, ref_image_path):
- 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)
- ref_image = cv2.imread(ref_image_path)
- 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)
- query_gray = cv2.cvtColor(query_image, cv2.COLOR_BGR2GRAY)
- ref_gray = cv2.cvtColor(ref_image, cv2.COLOR_BGR2GRAY)
- if os.path.exists(json_ref_path):
- ref_gray= process_image_with_mask(ref_image, scale_factor, json_ref_path,ref_gray)
- 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)
- # print(f"{filename}\n")
- # matched_query_kpts =[query_kpts[m.queryIdx].pt for m in matches]
- # matched_ref_kpts =[ref_kpts[m.trainIdx].pt for m in matches]
- # diff_query_ref=np.array(matched_ref_kpts) -np.array(matched_query_kpts)
- # # 计算每一列的模
- # diff_query_ref = np.linalg.norm(diff_query_ref, axis=0, ord=1)
- # # 计算所有行的平均值
- # diff_query_ref = np.mean(diff_query_ref)
- # # for index in range(len(matched_query_kpts)):
- # # diff_query_ref.append(matched_query_kpts[index]-matched_ref_kpts[index])
- #
- # print(matched_query_kpts)
- # print(matched_ref_kpts)
- # print(diff_query_ref)
- 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,
- )
- match_file_name = f"match_image_{filename}.jpg"
- cv2.imwrite(os.path.join(demo_image_path , match_file_name), matched_image)
- wrap_image = cv.warpPerspective(query_image_resize, M, (ref_image_resize.shape[1], ref_image_resize.shape[0]))
- # wrap_image = cv.warpPerspective(query_image, M,(ref_image.shape[1], ref_image.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)
- 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, result_file_name), result_image)
- return wraped_image_list
- def read_params_from_yml(yml_file_path):
- with open(yml_file_path, 'r') as file:
- params = yaml.safe_load(file)
- return params
- if __name__ == "__main__":
- yml_file_path = 'params.yml'
- params = read_params_from_yml(yml_file_path)
- registration_demo(params['image_dir'],params['demo_image_path'], params['json_ref_path'], params['ref_image_path'])
|