add: integrated SDK API

-add new test.py file which provide unifired
anomaly detection sdk in class formatwq

README.md

@@ -1,114 +1,18 @@
-<p align="center">
-<a href="https://github.com/xmba15/superpoint_superglue_deployment/actions/workflows/build.yml" target="_blank">
-  <img src="https://github.com/xmba15/superpoint_superglue_deployment/actions/workflows/build.yml/badge.svg" alt="Build Status">
-</a>
-</p>
+# Welding-Parts-Anomaly-Detection
+1 demo code
+python demo.py
 
-# 📝 simple library to make life easy when deploying superpoint, superglue models
+2 SDK code
+python test.py
 
----
+3 conda env install
+conda create -n superpoint_and_superglue python=3.8
+pip install -r requirements.txt
 
-## :gear: Installation
+4 data and code location
+url:192.168.20.250
+directory:/data/liudan/superpoint_superglue_deployment
+conda env name: superpoint_and_superglue
 
----
-
-```bash
-pip install superpoint_superglue_deployment
-```
-
-## :tada: TODO
-
----
-
-- [x] interface to deploy superpoint, superglue
-- [x] testing on real data
-
-## :running: How to Run
-
----
-
-### Basic usage
-
-```python
-import cv2
-import numpy as np
-from loguru import logger
-
-from superpoint_superglue_deployment import Matcher
-
-
-def main():
-    query_image = cv2.imread("./data/images/one_pillar_pagoda_1.jpg")
-    ref_image = cv2.imread("./data/images/one_pillar_pagoda_2.jpg")
-
-    query_gray = cv2.imread("./data/images/one_pillar_pagoda_1.jpg", 0)
-    ref_gray = cv2.imread("./data/images/one_pillar_pagoda_2.jpg", 0)
-
-    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,
-        query_kpts,
-        ref_image,
-        ref_kpts,
-        matches[:50],
-        None,
-        flags=2,
-    )
-    cv2.imwrite("matched_image.jpg", matched_image)
-
-
-if __name__ == "__main__":
-    main()
-```
-
-<p align="center">
-  <img src="https://raw.githubusercontent.com/xmba15/superpoint_superglue_deployment/master/docs/images/matched_image.jpg" alt="matched image sample">
-</p>
-
-- [Notebook with detailed sample code for SuperPoint](notebooks/demo_superpoint.ipynb)
-- [Notebook with detailed sample code for SuperGlue](notebooks/demo_superglue.ipynb)
-
-- Command line to test matching two images after installing the library
-
-```bash
-match_two_images --query_path [path/to/query/image] --ref_path [path/to/reference/image] --use_gpu
-```
-
-## 🎛 Development environment
-
----
-
-```bash
-mamba env create --file environment.yml
-mamba activate superpoint_superglue_deployment
-```
-
-## :gem: References
-
----
-
-- [SuperPoint: Self-Supervised Interest Point Detection and Description.](https://github.com/rpautrat/SuperPoint)
-- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://github.com/magicleap/SuperGluePretrainedNetwork)
+code maintainer:
+liudan, fengyang, zhengbochao

data/json/image165214-001.json

@@ -257,6 +257,86 @@
       "group_id": null,
       "shape_type": "rectangle",
       "flags": {}
+    },
+    {
+      "label": "D1",
+      "points": [
+        [
+          2392.077922077922,
+          700.0
+        ],
+        [
+          2411.5584415584417,
+          722.077922077922
+        ]
+      ],
+      "group_id": null,
+      "shape_type": "rectangle",
+      "flags": {}
+    },
+    {
+      "label": "D2",
+      "points": [
+        [
+          2421.9480519480517,
+          711.6883116883117
+        ],
+        [
+          2446.6233766233768,
+          736.3636363636364
+        ]
+      ],
+      "group_id": null,
+      "shape_type": "rectangle",
+      "flags": {}
+    },
+    {
+      "label": "D3",
+      "points": [
+        [
+          2394.6753246753246,
+          984.4155844155844
+        ],
+        [
+          2414.155844155844,
+          1006.4935064935065
+        ]
+      ],
+      "group_id": null,
+      "shape_type": "rectangle",
+      "flags": {}
+    },
+    {
+      "label": "D4",
+      "points": [
+        [
+          2423.246753246753,
+          966.2337662337662
+        ],
+        [
+          2450.5194805194806,
+          994.8051948051948
+        ]
+      ],
+      "group_id": null,
+      "shape_type": "rectangle",
+      "flags": {}
+    },
+    {
+      "label": "D5",
+      "points": [
+        [
+          2162.207792207792,
+          845.4545454545454
+        ],
+        [
+          2192.077922077922,
+          874.025974025974
+        ]
+      ],
+      "group_id": null,
+      "shape_type": "rectangle",
+      "flags": {}
     }
   ],
   "imagePath": "image165214-001.jpg",

data/json/image165214-019.json

@@ -3,102 +3,6 @@
   "flags": {},
   "shapes": [
     {
-      "label": "K",
-      "points": [
-        [
-          3513.90243902439,
-          467.80487804878044
-        ],
-        [
-          3570.0,
-          526.3414634146342
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "L",
-      "points": [
-        [
-          3416.341463414634,
-          470.2439024390244
-        ],
-        [
-          3477.3170731707314,
-          523.9024390243902
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "M",
-      "points": [
-        [
-          3021.2195121951218,
-          599.5121951219512
-        ],
-        [
-          3209.0243902439024,
-          701.951219512195
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "N",
-      "points": [
-        [
-          2518.780487804878,
-          328.7804878048781
-        ],
-        [
-          2699.2682926829266,
-          421.4634146341464
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "O",
-      "points": [
-        [
-          2235.8536585365855,
-          480.0
-        ],
-        [
-          2321.2195121951218,
-          582.4390243902438
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "P",
-      "points": [
-        [
-          2160.243902439024,
-          845.8536585365854
-        ],
-        [
-          2340.731707317073,
-          945.8536585365854
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
       "label": "Q",
       "points": [
         [
@@ -115,22 +19,6 @@
       "flags": {}
     },
     {
-      "label": "R",
-      "points": [
-        [
-          1043.1707317073171,
-          850.7317073170732
-        ],
-        [
-          1355.3658536585365,
-          1026.341463414634
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
       "label": "S",
       "points": [
         [
@@ -209,6 +97,38 @@
       "group_id": null,
       "shape_type": "rectangle",
       "flags": {}
+    },
+    {
+      "label": "R1",
+      "points": [
+        [
+          1126.0975609756097,
+          916.5853658536585
+        ],
+        [
+          1172.439024390244,
+          962.9268292682926
+        ]
+      ],
+      "group_id": null,
+      "shape_type": "rectangle",
+      "flags": {}
+    },
+    {
+      "label": "R2",
+      "points": [
+        [
+          1177.3170731707316,
+          897.0731707317073
+        ],
+        [
+          1223.658536585366,
+          948.2926829268293
+        ]
+      ],
+      "group_id": null,
+      "shape_type": "rectangle",
+      "flags": {}
     }
   ],
   "imagePath": "image165214-019.jpg",

data/json/image165214-037.json

@@ -3,70 +3,6 @@
   "flags": {},
   "shapes": [
     {
-      "label": "T",
-      "points": [
-        [
-          3635.853658536585,
-          838.5365853658536
-        ],
-        [
-          3828.5365853658536,
-          953.170731707317
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "U",
-      "points": [
-        [
-          3484.6341463414633,
-          1050.731707317073
-        ],
-        [
-          3716.341463414634,
-          1419.0243902439022
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "V",
-      "points": [
-        [
-          3355.3658536585363,
-          1048.292682926829
-        ],
-        [
-          3455.3658536585363,
-          1214.1463414634145
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
-      "label": "W",
-      "points": [
-        [
-          2960.243902439024,
-          921.4634146341463
-        ],
-        [
-          3130.9756097560976,
-          1023.9024390243902
-        ]
-      ],
-      "group_id": null,
-      "shape_type": "rectangle",
-      "flags": {}
-    },
-    {
       "label": "X",
       "points": [
         [

demo_env.py

@@ -35,7 +35,7 @@ def process_image_with_mask(ref_image, scale_factor, json_ref_path, ref_gray):
     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)
+    # cv2.imwrite('1111111.jpg',ref_gray_masked)
     return ref_gray_masked
 
 
@@ -43,87 +43,89 @@ 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格式）
-        try:
-            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)
-                # 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,
-                    query_kpts,
-                    ref_image,
-                    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)
-        except:
-            pass
+
+        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
 

generate_test_images.py

@@ -8,6 +8,7 @@ import os
 import cv2
 import json
 import numpy as np
+from loguru import logger
 from superpoint_superglue_deployment import Matcher
 
 class GenerateTestImage:
@@ -27,12 +28,12 @@ class GenerateTestImage:
             os.mkdir(self.save_dir)
 
 
-        self.template_img_paths = [os.path.join(self.template_img_dir, f) for f in os.listdir(self.template_img_dir) if
-                                   f.endswith('.jpg') or f.endswith('.JPG') ]
-        self.template_mask_json_paths = [os.path.join(self.template_mask_json_dir, f) for f in os.listdir(self.template_mask_json_dir) if
-                                         f.endswith('.json')]
-        self.template_label_json_paths = [os.path.join(self.template_label_json_dir, f) for f in os.listdir(self.template_label_json_dir) if
-                                          f.endswith('.json')]
+        self.template_img_paths = sorted([os.path.join(self.template_img_dir, f) for f in os.listdir(self.template_img_dir) if
+                                   f.endswith('.jpg') or f.endswith('.JPG') ])
+        self.template_mask_json_paths = sorted([os.path.join(self.template_mask_json_dir, f) for f in os.listdir(self.template_mask_json_dir) if
+                                         f.endswith('.json')])
+        self.template_label_json_paths = sorted([os.path.join(self.template_label_json_dir, f) for f in os.listdir(self.template_label_json_dir) if
+                                          f.endswith('.json')])
 
     def generate_test_image(self, query_image_dir):
         for file_name in os.listdir(query_image_dir):
@@ -165,11 +166,12 @@ class GenerateTestImage:
         return best_match_index
 
 if __name__ =="__main__":
-    template_img_dir = './data/REF'
-    template_mask_json_dir = './data/JSON_MASK/'
-    template_label_json_dir = './data/JSON/'
-    save_dir = "./data/generated"
+    template_img_dir = './data/ref/'
+    template_mask_json_dir = './data/mask_json/'
+    template_label_json_dir = './data/json/'
+    save_dir = "./data/D_result"
 
-    query_image_dir = "/data2/liudan/superpoint_superglue_deployment/data/QUERY"
+    query_image_dir = "./data/query"
+    # match_path = './data/vague_match'
     gti =GenerateTestImage(template_img_dir, template_mask_json_dir, template_label_json_dir, save_dir)
     gti.generate_test_image(query_image_dir)

image_registration.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2024/6/26 0026 下午 3:21
+# @Author  : liudan
+# @File    : image_registration.py
+# @Software: pycharm
+import os
+import cv2
+import json
+import numpy as np
+from loguru import logger
+from superpoint_superglue_deployment import Matcher
+
+class GenerateTestImage:
+    def __init__(self, template_img_dir, template_mask_json_dir, save_dir):
+
+        self.template_img_dir = template_img_dir
+        self.template_mask_json_dir = template_mask_json_dir
+        self.save_dir = save_dir
+        if not os.path.exists(self.save_dir):
+            os.mkdir(self.save_dir)
+
+
+        self.template_img_paths = sorted([os.path.join(self.template_img_dir, f) for f in os.listdir(self.template_img_dir) if
+                                   f.endswith('.jpg') or f.endswith('.JPG') ])
+        self.template_mask_json_paths = sorted([os.path.join(self.template_mask_json_dir, f) for f in os.listdir(self.template_mask_json_dir) if
+                                         f.endswith('.json')])
+
+    def generate_test_image(self, query_image_dir):
+        for file_name in os.listdir(query_image_dir):
+            query_image_path = os.path.join(query_image_dir, file_name)
+            best_matched_index = self.find_best_matched_template_image(query_image_path)
+            best_matched_template_img_path = self.template_img_paths[best_matched_index]
+            best_matched_template_mask_json_path = self.template_mask_json_paths[best_matched_index]
+
+            query_image = cv2.imread(query_image_path)
+            ref_image = cv2.imread(best_matched_template_img_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)
+            if os.path.exists(best_matched_template_mask_json_path):
+                ref_gray = self.process_image_with_mask(ref_image, scale_factor, best_matched_template_mask_json_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)
+            try:
+                if len(matches) < 4:
+                    raise ValueError("Not enough matches to compute homography.")
+                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,
+                )
+                if M is None:
+                    raise ValueError("Unable to compute homography.")
+
+                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)
+
+                wrap_image = cv2.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]))
+                sub_image = cv2.subtract(ref_image, wrap_image)
+
+                cv2.imwrite(os.path.join(self.save_dir, f'match_{file_name[:-4]}.jpg'), matched_image)
+                cv2.imwrite(os.path.join(self.save_dir, f'wrap_{file_name[:-4]}.jpg'), wrap_image)
+                cv2.imwrite(os.path.join(self.save_dir, f'sub_{file_name[:-4]}.jpg'), sub_image)
+
+            except ValueError as e:
+                print(e)
+
+            except Exception as e:
+                print(f"An error occurred: {e}")
+
+
+    def process_image_with_mask(self, 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)
+        return ref_gray_masked
+
+    def find_best_matched_template_image(self, query_img_path):
+        number_match_pts_list = []
+        for index, item in enumerate(self.template_img_paths):
+            ref_image_path = item
+            query_image = cv2.imread(query_img_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)
+            if os.path.exists( self.template_mask_json_paths[index]):
+                ref_gray = self.process_image_with_mask(ref_image, scale_factor,self.template_mask_json_paths[index],
+                                                        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)
+            number_match_pts_list.append(len(matches))
+        best_match_index = number_match_pts_list.index(max(number_match_pts_list))
+        return best_match_index
+
+if __name__ =="__main__":
+    template_img_dir = './data/best_registration/template_image/'
+    template_mask_json_dir = './data/best_registration/mask_json/'
+    save_dir = "./data/best_registration/result"
+
+    query_image_dir = "./data/best_registration/query_image"
+    gti =GenerateTestImage(template_img_dir, template_mask_json_dir, save_dir)
+    gti.generate_test_image(query_image_dir)

num_pixel.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2024/6/20 0020 下午 1:44
+# @Author  : liudan
+# @File    : 11111.py
+# @Software: pycharm
+
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2024/6/20 0020 下午 1:44
+# @Author  : liudan
+# @File    : 11111.py
+# @Software: pycharm
+
+import os
+import cv2
+import cv2 as cv
+import numpy as np
+import json
+import os
+
+
+THRESHOLD=0.15
+
+def otsu_binarize(image):
+
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    gray_image = cv2.equalizeHist(gray_image)
+    # ret1, mask1 = cv2.threshold(gray_image, 0, 255, cv2.THRESH_BINARY +cv2.THRESH_OTSU)
+    ret2, mask =cv2.threshold(gray_image,60, 255, cv2.THRESH_BINARY)
+    # mask =mask1 & mask2
+    kernel = np.ones((3, 3), np.uint8)
+    mask = cv2.dilate(mask, kernel, iterations=3)
+    mask = cv2.erode(mask, kernel, iterations=3)
+
+
+    return mask
+
+def read_json_file(json_path):
+    with open(json_path, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+
+    coordinates_dict = {}
+    for shape in data['shapes']:
+        shape_name = shape['label']
+        points = shape['points']
+        rounded_points = [[int(round(x)), int(round(y))] for x, y in points]
+        coordinates_dict[shape_name] = rounded_points
+    return coordinates_dict
+
+
+
+def pixel_calculation(ref_image_dir, query_positive_image_dir, query_negative_image_dir, json_path):
+    true_positives = 0
+    false_positives = 0
+    false_negatives = 0
+    true_negatives = 0
+
+    false_negatives_list = []
+    false_positives_list = []
+
+
+    for ref_file in os.listdir(ref_image_dir):
+        ref_image_path = os.path.join(ref_image_dir, ref_file)
+        ref_image = cv2.imread(ref_image_path)
+        ostu_ref_image = otsu_binarize(ref_image)
+
+
+        for query_file in os.listdir(query_positive_image_dir):
+            query_image_path = os.path.join(query_positive_image_dir, query_file)
+            query_image = cv2.imread(query_image_path)
+            ostu_query_image = otsu_binarize(query_image)
+
+            coordinates = read_json_file(json_path)
+
+            num_pixel_list = []
+            for _, points in coordinates.items():
+                x1, y1 = points[0]
+                x2, y2 = points[1]
+
+                region_ref = ostu_ref_image[y1:y2, x1:x2]
+                region_query = ostu_query_image[y1:y2, x1:x2]
+
+                num_pixel_region_ref = round((np.sum(region_ref == 0) / (region_ref.shape[0] * region_ref.shape[1])),2)
+                num_pixel_region_query = round((np.sum(region_query == 0) / (region_query.shape[0] * region_query.shape[1])),2)
+
+                num_pixel_list.append(num_pixel_region_query)
+
+            flag = any(element <= THRESHOLD for element in num_pixel_list)
+
+            if not flag:
+                true_positives += 1
+            else:
+                false_negatives += 1
+                false_negatives_list.append((query_file, ref_file))
+                # cv2.imwrite(os.path.join(query_positive_image_dir, f'{query_file}_img.jpg'),ostu_query_image)
+                print(f"positive ={num_pixel_list}")
+
+
+
+        for query_file in os.listdir(query_negative_image_dir):
+            query_image_path = os.path.join(query_negative_image_dir, query_file)
+            query_image = cv2.imread(query_image_path)
+            ostu_query_image = otsu_binarize(query_image)
+            # cv2.imwrite(os.path.join(query_negative_image_dir, '1.jpg'), ostu_query_image)
+
+            coordinates = read_json_file(json_path)
+
+            num_pixel_list = []
+            for _, points in coordinates.items():
+                x1, y1 = points[0]
+                x2, y2 = points[1]
+
+                region_ref = ostu_ref_image[y1:y2, x1:x2]
+                region_query = ostu_query_image[y1:y2, x1:x2]
+
+                num_pixel_region_ref = round((np.sum(region_ref == 0) / (region_ref.shape[0] * region_ref.shape[1])), 2)
+                num_pixel_region_query = round((np.sum(region_query == 0) / (region_query.shape[0] * region_query.shape[1])), 2)
+                num_pixel_list.append(num_pixel_region_query)
+
+            flag = any(element <= THRESHOLD for element in num_pixel_list)
+
+            if flag:
+                true_negatives += 1
+            else:
+                false_positives += 1
+                false_positives_list.append((query_file, ref_file))
+                cv2.imwrite(os.path.join(query_negative_image_dir, f'{query_file}_img.jpg'), ostu_query_image)
+                print(f"negtive ={num_pixel_list}")
+
+    Accurary = (true_positives + true_negatives) / (true_positives + true_negatives + false_positives + false_negatives)
+    # Precision = true_positives / (true_positives + false_positives)
+    Recall = true_negatives / (true_negatives + false_positives)
+    Precision = true_negatives / (true_negatives + false_negatives)
+    # Recall = true_positives / (true_positives + false_negatives)
+
+    F1_score = 2 * (Precision * Recall) / (Precision + Recall)
+
+    print(false_negatives_list)
+    print(false_positives_list)
+    # print(file_name)
+    print(f"Accurary:{Accurary: .4f}")
+    print(f"Precision: {Precision:.4f}")
+    print(f"Recall: {Recall:.4f}")
+    print(f"F1 Score: {F1_score:.4f}")
+
+
+if __name__ == "__main__":
+    # ref_image_dir = 'H:\\dataset\\region_samples\\positive_test_img\\R_ref'
+    # query_positive_image_dir = 'H:\\dataset\\region_samples\\positive_test_img\\R_pos'
+    # query_negative_image_dir = 'H:\\dataset\\region_samples\\positive_test_img\\R_neg'
+    # json_path = 'H:\\dataset\\region_samples\\positive_test_img\\R_json\\2.json'
+    ref_image_dir = 'H:\\dataset\\A_ref'
+    query_positive_image_dir = 'H:\\dataset\\A'
+    query_negative_image_dir = 'H:\\dataset\\C'
+    json_path = 'H:\\dataset\\B\\image-024_R.json'
+    pixel_calculation(ref_image_dir, query_positive_image_dir, query_negative_image_dir, json_path)

requirements.txt

@@ -1,6 +1,45 @@
-# typing
-typing-extensions
-
-torch>=1.1.0
-opencv-python>=3.4.19
-loguru>=0.7.0
+certifi==2024.2.2
+charset-normalizer==3.3.2
+filelock==3.14.0
+fsspec==2024.5.0
+idna==3.7
+imageio==2.34.1
+Jinja2==3.1.4
+joblib==1.4.2
+lazy_loader==0.4
+loguru==0.7.2
+MarkupSafe==2.1.5
+mpmath==1.3.0
+networkx==3.1
+numpy==1.24.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.5.40
+nvidia-nvtx-cu12==12.1.105
+opencv-python==4.9.0.80
+packaging==24.0
+pillow==10.3.0
+PyWavelets==1.4.1
+PyYAML==6.0.1
+requests==2.32.2
+scikit-image==0.21.0
+scikit-learn==1.3.2
+scipy==1.10.1
+superpoint-superglue-deployment==0.0.3
+sympy==1.12
+threadpoolctl==3.5.0
+tifffile==2023.7.10
+torch==1.12.0+cu116
+torchaudio==0.12.0+rocm5.1.1
+torchvision==0.13.0+cu116
+triton==2.3.0
+typing_extensions==4.11.0
+urllib3==2.2.1

rorate_image_detection.py

@@ -20,12 +20,14 @@ import cv2
 import numpy as np
 import random
 
+
+# 在模板图的基础上创造负样本（翻转图片）
 def rotate_image(image_path):
     rotate_result_images = []
     for path in os.listdir(image_path):
         ref_image_path = os.path.join(image_path,path)
         ref_image = cv2.imread(ref_image_path, 0)
-        if random.random() < 0.5:
+        if random.random() < 1:  # 以0.5的概率决定是否翻转
             rotated_image = cv2.rotate(ref_image, cv2.ROTATE_180)
             rotate_result_images.append((rotated_image, path))
         else:
@@ -35,6 +37,7 @@ def rotate_image(image_path):
     return rotate_result_images
 
 
+# 配准操作，计算模板图与查询图配准点的距离差异（翻转后的距离大于未翻转的）
 def match_image(ref_img, query_img):
     superglue_matcher = Matcher(
         {
@@ -48,45 +51,54 @@ def match_image(ref_img, query_img):
             "use_gpu": True,
         }
     )
-    _, _, _, _, matches = superglue_matcher.match(ref_img, query_img)
-    return matches
-
+    query_kpts, ref_kpts, _, _, matches = superglue_matcher.match(ref_img, query_img)
+    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)
+    print(diff_query_ref)
+    if len(matches)!=0:
+        diff_query_ref = np.linalg.norm(diff_query_ref, axis=1, ord=2)
+        diff_query_ref = np.sqrt(diff_query_ref)
+        diff_query_ref = np.mean(diff_query_ref)
+    else:
+        diff_query_ref = np.inf
+    return diff_query_ref  # 返回差异值，用于后续做比较
 
 
 
 
 def registration(ref_image_dir, ref_selection_dir):
-    # filename = os.listdir(ref_image_dir)
-    # image_files = [f for f in filename if f.endswith(('.jpg', '.jpeg', '.png', '.JPG'))]
-    # image_file = random.choice(image_files)
-    # ref_image_path = os.path.join(ref_image_dir, image_file)
-    # ref_image = cv2.imread(ref_image_path, 0)
 
+    # 计算recall等需要的的参数
     true_positives = 0
     false_positives = 0
     false_negatives = 0
     true_negatives = 0
 
+    # 返回预测错误的样本
     positive_false = []
     negative_false = []
 
+    # 循环遍历模板图
     for file_name in os.listdir(ref_selection_dir):
         ref_image_path = os.path.join(ref_selection_dir,file_name)
         ref_image = cv2.imread(ref_image_path, 0)
 
+        save_path = './data/region_rotate/vague/A' # 保存预测错误的原图片，可不用
 
-        # positive_false = []
-        save_path = './data/region_rotate/A'
+        # 循环遍历正样本（这里正样本是查询图，因为查询图都是没有翻转的，所以可当作正样本）
+        # 这是在正样本中预测
         for file in os.listdir(ref_image_dir):
             query_image_path = os.path.join(ref_image_dir, file)
             query_image = cv2.imread(query_image_path, 0)
             query_image_rotate = cv2.rotate(query_image, cv2.ROTATE_180)
             # cv2.imwrite('1111111111.jpg', query_image_rotate)
 
-            matches1 = match_image(ref_image, query_image)
-            matches2 = match_image(ref_image,query_image_rotate)
+            diff1 = match_image(ref_image, query_image) # 计算模板图与查询图的配准点差异
+            diff2 = match_image(ref_image,query_image_rotate) # 计算模板图与翻转180度图的配准点差异
 
-            if len(matches1) > len(matches2):
+            # if (len(matches1) > len(matches2)) or (diff1<diff2) :
+            if diff1 < diff2:
                 flag = True
             else:
                 flag = False
@@ -98,10 +110,10 @@ def registration(ref_image_dir, ref_selection_dir):
                 false_negatives += 1
                 positive_false.append((file,file_name))
                 cv2.imwrite(os.path.join(save_path, f'{file[:-4]}.jpg'),query_image)
-        # print(positive_false)
-        # print(file_name)
 
-        # negative_false = []
+
+        # 这是在负样本中预测
+        # 因为缺少负样本，所以用rotate_image方法创造负样本
         nagetive_image = rotate_image(ref_image_dir)
         for i, item in enumerate(nagetive_image):
             file, path = item
@@ -109,10 +121,11 @@ def registration(ref_image_dir, ref_selection_dir):
             query_image = file
             query_image_rotate = cv2.rotate(query_image, cv2.ROTATE_180)
 
-            matches1 = match_image(ref_image, query_image)
-            matches2 = match_image(ref_image, query_image_rotate)
+            diff1 = match_image(ref_image, query_image)
+            diff2 = match_image(ref_image, query_image_rotate)
 
-            if len(matches1) > len(matches2):
+            # if (len(matches1) > len(matches2)) or (diff1<diff2):
+            if diff1 < diff2:
                 flag = True
             else:
                 flag = False
@@ -124,8 +137,6 @@ def registration(ref_image_dir, ref_selection_dir):
                 false_positives += 1
                 negative_false.append((path, file_name))
                 cv2.imwrite(os.path.join(save_path, f'{path[:-4]}.jpg'),query_image)
-        # print(negative_false)
-        # print(file_name)
 
     Accurary = (true_positives + true_negatives) / (true_positives + true_negatives + false_positives + false_negatives)
     Precision = true_negatives / (true_negatives + false_negatives)
@@ -145,8 +156,10 @@ def registration(ref_image_dir, ref_selection_dir):
 
 
 if __name__ == "__main__":
-    ref_image_path = './data/region_rotate/E1'
-    ref_selection_dir = './data/region_rotate/E'
+    # ref_image_path = './data/region_rotate/vague/E2'
+    # ref_selection_dir = './data/region_rotate/vague/E22'
+    ref_image_path = './data/big_624/G1'
+    ref_selection_dir = './data/big_624/G2'
     # start_time = time.time()
     registration(ref_image_path, ref_selection_dir)
     # end_time = time.time()

rotate_env.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2024/6/24 0024 上午 10:38
+# @Author  : liudan
+# @File    : rotate_env.py
+# @Software: pycharm
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2024/6/19 0019 下午 4:35
+# @Author  : liudan
+# @File    : rorate_image_detection.py
+# @Software: pycharm
+import time
+
+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
+
+import cv2
+import numpy as np
+import random
+
+
+# 在模板图的基础上创造负样本（翻转图片）
+def rotate_image(image_path):
+    rotate_result_images = []
+    for path in os.listdir(image_path):
+        ref_image_path = os.path.join(image_path,path)
+        ref_image = cv2.imread(ref_image_path, 0)
+        if random.random() < 0.5:  # 以0.5的概率决定是否翻转
+            rotated_image = cv2.rotate(ref_image, cv2.ROTATE_180)
+            rotate_result_images.append((rotated_image, path))
+        else:
+            print(f"Skipped")
+    print(len(rotate_result_images))
+
+    return rotate_result_images
+
+
+# 配准操作，计算模板图与查询图配准点的距离差异（翻转后的距离大于未翻转的）
+def match_image(ref_img, query_img):
+    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(ref_img, query_img)
+    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)
+    print(diff_query_ref)
+    if len(matches)!=0:
+        diff_query_ref = np.linalg.norm(diff_query_ref, axis=1, ord=2)
+        diff_query_ref = np.sqrt(diff_query_ref)
+        diff_query_ref = np.mean(diff_query_ref)
+    else:
+        diff_query_ref = np.inf
+    return diff_query_ref, matches  # 返回差异值，用于后续做比较
+
+
+
+
+
+def registration(ref_image_dir, ref_selection_dir):
+
+    # 计算recall等需要的的参数
+    true_positives = 0
+    false_positives = 0
+    false_negatives = 0
+    true_negatives = 0
+
+    # 返回预测错误的样本
+    positive_false = []
+    negative_false = []
+
+    # 循环遍历模板图
+    for file_name in os.listdir(ref_selection_dir):
+        ref_image_path = os.path.join(ref_selection_dir,file_name)
+        ref_image = cv2.imread(ref_image_path, 0)
+
+        save_path = './data/region_rotate/vague/A' # 保存预测错误的原图片，可不用
+
+        # 循环遍历正样本（这里正样本是查询图，因为查询图都是没有翻转的，所以可当作正样本）
+        # 这是在正样本中预测
+        for file in os.listdir(ref_image_dir):
+            query_image_path = os.path.join(ref_image_dir, file)
+            query_image = cv2.imread(query_image_path, 0)
+            query_image_rotate = cv2.rotate(query_image, cv2.ROTATE_180)
+            # cv2.imwrite('1111111111.jpg', query_image_rotate)
+
+            diff1, matches1 = match_image(ref_image, query_image) # 计算模板图与查询图的配准点差异
+            diff2, matches2 = match_image(ref_image,query_image_rotate) # 计算模板图与翻转180度图的配准点差异
+
+            # if (len(matches1) > len(matches2)) or (diff1<diff2) :
+            if diff1 < diff2:
+                flag = True
+            else:
+                flag = False
+            print(flag)
+
+            if flag == True:
+                true_positives += 1
+            else:
+                false_negatives += 1
+                positive_false.append((file,file_name))
+                cv2.imwrite(os.path.join(save_path, f'{file[:-4]}.jpg'),query_image)
+
+
+        # 这是在负样本中预测
+        # 因为缺少负样本，所以用rotate_image方法创造负样本
+        nagetive_image = rotate_image(ref_image_dir)
+        for i, item in enumerate(nagetive_image):
+            file, path = item
+            # query_image = cv2.cvtColor(file,cv2.COLOR_BGR2GRAY)
+            query_image = file
+            query_image_rotate = cv2.rotate(query_image, cv2.ROTATE_180)
+
+            diff1, matches1 = match_image(ref_image, query_image)
+            diff2, matches2 = match_image(ref_image, query_image_rotate)
+
+            # if (len(matches1) > len(matches2)) or (diff1<diff2):
+            if diff1 < diff2:
+                flag = True
+            else:
+                flag = False
+            print(flag)
+
+            if flag == False:
+                true_negatives += 1
+            else:
+                false_positives += 1
+                negative_false.append((path, file_name))
+                cv2.imwrite(os.path.join(save_path, f'{path[:-4]}.jpg'),query_image)
+
+    Accurary = (true_positives + true_negatives) / (true_positives + true_negatives + false_positives + false_negatives)
+    Precision = true_negatives / (true_negatives + false_negatives)
+    Recall = true_negatives / (true_negatives + false_positives)
+    F1_score = 2 * (Precision * Recall) / (Precision + Recall)
+
+    print(positive_false)
+    print(negative_false)
+    # print(file_name)
+    print(f"Accurary:{Accurary: .4f}")
+    print(f"Precision: {Precision:.4f}")
+    print(f"Recall: {Recall:.4f}")
+    print(f"F1 Score: {F1_score:.4f}")
+
+
+
+
+
+if __name__ == "__main__":
+    ref_image_path = './data/region_rotate/vague/E2'
+    ref_selection_dir = './data/region_rotate/vague/E22'
+    # start_time = time.time()
+    registration(ref_image_path, ref_selection_dir)
+    # end_time = time.time()
+    # elapsed_time = end_time - start_time
+    # print(f"程序用时: {elapsed_time:.2f} 秒")
+

superpoint_superglue_deployment/README.md

@@ -0,0 +1,114 @@
+<p align="center">
+<a href="https://github.com/xmba15/superpoint_superglue_deployment/actions/workflows/build.yml" target="_blank">
+  <img src="https://github.com/xmba15/superpoint_superglue_deployment/actions/workflows/build.yml/badge.svg" alt="Build Status">
+</a>
+</p>
+
+# 📝 simple library to make life easy when deploying superpoint, superglue models
+
+---
+
+## :gear: Installation
+
+---
+
+```bash
+pip install superpoint_superglue_deployment
+```
+
+## :tada: TODO
+
+---
+
+- [x] interface to deploy superpoint, superglue
+- [x] testing on real data
+
+## :running: How to Run
+
+---
+
+### Basic usage
+
+```python
+import cv2
+import numpy as np
+from loguru import logger
+
+from superpoint_superglue_deployment import Matcher
+
+
+def main():
+    query_image = cv2.imread("./data/images/one_pillar_pagoda_1.jpg")
+    ref_image = cv2.imread("./data/images/one_pillar_pagoda_2.jpg")
+
+    query_gray = cv2.imread("./data/images/one_pillar_pagoda_1.jpg", 0)
+    ref_gray = cv2.imread("./data/images/one_pillar_pagoda_2.jpg", 0)
+
+    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,
+        query_kpts,
+        ref_image,
+        ref_kpts,
+        matches[:50],
+        None,
+        flags=2,
+    )
+    cv2.imwrite("matched_image.jpg", matched_image)
+
+
+if __name__ == "__main__":
+    main()
+```
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/xmba15/superpoint_superglue_deployment/master/docs/images/matched_image.jpg" alt="matched image sample">
+</p>
+
+- [Notebook with detailed sample code for SuperPoint](notebooks/demo_superpoint.ipynb)
+- [Notebook with detailed sample code for SuperGlue](notebooks/demo_superglue.ipynb)
+
+- Command line to test matching two images after installing the library
+
+```bash
+match_two_images --query_path [path/to/query/image] --ref_path [path/to/reference/image] --use_gpu
+```
+
+## 🎛 Development environment
+
+---
+
+```bash
+mamba env create --file environment.yml
+mamba activate superpoint_superglue_deployment
+```
+
+## :gem: References
+
+---
+
+- [SuperPoint: Self-Supervised Interest Point Detection and Description.](https://github.com/rpautrat/SuperPoint)
+- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://github.com/magicleap/SuperGluePretrainedNetwork)

superpoint_superglue_deployment/matcher.py

@@ -32,14 +32,14 @@ class Matcher:
 
     def __init__(
         self,
-        config: Optional[Dict[str, Any]] = None,
+        config: Optional[Dict[str, Any]] = None, shape=160
     ):
         self._config = self.__DEFAULT_CONFIG.copy()
         if config is not None:
             self._config.update(config)
         self._config["superpoint"].update({"use_gpu": self._config["use_gpu"]})
         self._config["superglue"].update({"use_gpu": self._config["use_gpu"]})
-        self._superpoint_handler = SuperPointHandler(self._config["superpoint"])
+        self._superpoint_handler = SuperPointHandler(self._config["superpoint"],shape)
         self._superglue_handler = SuperGlueHandler(self._config["superglue"])
 
     def match(

superpoint_superglue_deployment/superpoint_handler.py

@@ -31,8 +31,9 @@ class SuperPointHandler:
 
     def __init__(
         self,
-        config: Optional[Dict[str, Any]] = None,
+        config: Optional[Dict[str, Any]] = None, shape=160
     ):
+        self.shape = shape
         self._config = self.__DEFAULT_CONFIG.copy()
         if config is not None:
             self._config.update(config)
@@ -59,7 +60,7 @@ class SuperPointHandler:
 
     def _validate_input_shape(self, image_shape: Tuple[int, int]):
         assert (
-            max(image_shape) >= 160 and max(image_shape) <= 4100
+            max(image_shape) >= self.shape and max(image_shape) <= 4100
         ), f"input resolution {image_shape} is too small or too large"
 
     @property

test.py

@@ -16,23 +16,41 @@ from superpoint_superglue_deployment import Matcher
 import numpy as np
 from loguru import logger
 import cv2 as cv
+import yaml
 
 
-# 定义大类
 class StructureClass:
-    def __init__(self, ref_image_path, query_image_path, json_path, save_image_path, json_mask_path):
+    def __init__(self, ref_image_path, query_image_path, json_path, save_image_path, json_mask_path, scale_factor=0.37):
         self.weldmentclasses = []
         self.ref_image = []
         self.query_image = []
         self.boxes_xy_label = {}
-        self.scale_factor = 0.37
+        self.scale_factor = scale_factor
+        self.superglue_matcher = Matcher({
+            "superpoint": {
+                "input_shape": (-1, -1),
+                "keypoint_threshold": 0.003,
+            },
+            "superglue": {
+                "match_threshold": 0.5,
+            },
+            "use_gpu": True,
+        })
+
 
 
         self.read_ref_image(ref_image_path)
         self.read_query_image(query_image_path)
         self.read_json(json_path)
 
-        self.registration_demo(save_image_path, json_mask_path)
+        _, wrap_image, _ = self.registration_demo(save_image_path, json_mask_path)
+        self.replace_query_image(wrap_image)
+        self.group_weldment()
+
+    # 定义一个函数来获取每个元素的首字母
+    def get_first_letter(self, item):
+        it = '-'.join(item.split('-')[:2])
+        return it
 
     def read_ref_image(self, path):
         self.ref_image = self.process_image_data(path)
@@ -44,6 +62,13 @@ class StructureClass:
         self.query_image = wrap_image
 
     def read_json(self, json_path):
+        """
+        读取焊接件的标注信息
+        :param json_path:
+        :type json_path:
+        :return:
+        :rtype:
+        """
         with open(json_path, 'r') as f:
             data = json.load(f)
         for shape in data['shapes']:
@@ -55,6 +80,13 @@ class StructureClass:
                 self.boxes_xy_label[label] = [x1, y1, x2, y2]
 
     def process_image_data(self, data):
+        """
+        读取图像，如果是文件则用cv读取，如果已经读取过则直接返回
+        :param data:
+        :type data:
+        :return:
+        :rtype:
+        """
         extensions = ['.PNG', '.png', '.jpg', '.jpeg', '.JPG', '.JPEG']
 
         if any(data.endswith(ext) for ext in extensions):
@@ -68,96 +100,71 @@ class StructureClass:
             else:
                 raise FileNotFoundError(f"Image file {data} not found")
 
-    def registration_demo(self,save_image_path, json_mask_path):
-        height, width = self.ref_image.shape[:2]
-        ref_image_resize = cv2.resize(self.ref_image, dsize=None, fx=self.scale_factor, fy=self.scale_factor)
-        query_image_resize = cv2.resize(self.query_image, dsize=None, fx=self.scale_factor, fy=self.scale_factor)
-        ref_gray = cv2.cvtColor(ref_image_resize, cv2.COLOR_BGR2GRAY)
-        query_gray = cv2.cvtColor(query_image_resize, cv2.COLOR_BGR2GRAY)
+    def mask2image(self, json_mask_path, ref_gray):
+        """
+        如果image_mask json文件存在，则将ref_image按照mask的形状扣出来
+        :param json_mask_path:
+        :type json_mask_path:
+        :param ref_gray:
+        :type ref_gray:
+        :return:
+        :rtype:
+        """
+        with open(json_mask_path, 'r') as f:
+            data = json.load(f)
+        shape = data['shapes'][0]
+        if shape['shape_type'] == 'polygon':
+            coords = [(int(round(x * self.scale_factor)), int(round(y * self.scale_factor))) for x, y in
+                      shape['points']]
+            mask = np.zeros(ref_gray.shape, np.uint8)
+            cv2.fillPoly(mask, [np.array(coords, np.int32)], 255)
+            ref_gray = cv2.bitwise_and(ref_gray, mask)
 
-        if os.path.exists(json_path):
-            with open(json_mask_path, 'r') as f:
-                data = json.load(f)
-            shapes = data['shapes']
-            shape = shapes[0]
-            if shape['shape_type'] == 'polygon':
-                coords = [(int(round(x * self.scale_factor)), int(round(y * self.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 = cv2.bitwise_and(ref_gray, ref_gray, mask=mask)
+        return 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)
-        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.jpg"
-        cv2.imwrite(os.path.join(save_image_path, match_file_name), matched_image)
+    def registration_demo(self, save_image_path, json_mask_path):
+        ref_image_resize = cv2.resize(self.ref_image, None, fx=self.scale_factor, fy=self.scale_factor)
+        query_image_resize = cv2.resize(self.query_image, None, fx=self.scale_factor, fy=self.scale_factor)
+        ref_gray = cv2.cvtColor(ref_image_resize, cv2.COLOR_BGR2GRAY)
+        query_gray = cv2.cvtColor(query_image_resize, cv2.COLOR_BGR2GRAY)
+
+        if os.path.exists(json_mask_path):
+            ref_gray = self.mask2image(json_mask_path, ref_gray)
 
-        wrap_file_name = f"wrap.jpg"
-        wrap_image = cv.warpPerspective(query_image_resize, M, (ref_image_resize.shape[1], ref_image_resize.shape[0]))
-        wrap_image = cv2.resize(wrap_image, (self.ref_image.shape[1], self.ref_image.shape[0]))
-        cv2.imwrite(os.path.join(save_image_path, wrap_file_name), wrap_image)
+        query_kpts, ref_kpts, _, _, matches = self.superglue_matcher.match(query_gray, ref_gray)
+        src_pts = np.float32([query_kpts[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2)
+        dst_pts = np.float32([ref_kpts[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2)
+        M, mask = cv2.findHomography(src_pts, dst_pts, cv2.USAC_MAGSAC, 5.0, maxIters=10000, confidence=0.95)
+        logger.info(f"Number of inliers: {mask.sum()}")
+        matches = [m for m, m_mask in zip(matches, mask) if m_mask]
+        matches.sort(key=lambda m: m.distance)
+        matched_image = cv2.drawMatches(query_image_resize, query_kpts, ref_image_resize, ref_kpts, matches[:50], None, flags=2)
 
-        sub_file_name = f"result.jpg"
+        wrap_image = cv2.warpPerspective(query_image_resize, M, (ref_image_resize.shape[1], ref_image_resize.shape[0]))
+        wrap_image = cv2.resize(wrap_image, (self.ref_image.shape[1], self.ref_image.shape[0]))
         sub_image = cv2.subtract(self.ref_image, wrap_image)
-        cv2.imwrite(os.path.join(save_image_path, sub_file_name), sub_image)
 
-        return matched_image, wrap_image, sub_image
+        cv2.imwrite(os.path.join(save_image_path, "match.jpg"), matched_image)
+        cv2.imwrite(os.path.join(save_image_path, "wrap.jpg"), wrap_image)
+        cv2.imwrite(os.path.join(save_image_path, "result.jpg"), sub_image)
 
+        return matched_image, wrap_image, sub_image
 
-    # def process_image_with_mask(self, json_mask_path, ref_gray):
-    #     with open(json_mask_path, 'r') as f:
-    #         data = json.load(f)
-    #     shapes = data['shapes']
-    #     shape = shapes[0]
-    #     if shape['shape_type'] == 'polygon':
-    #         coords = [(int(round(x * self.scale_factor)), int(round(y * self.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('ref_gray_mask.jpg', ref_gray_masked)
-    #     return ref_gray_masked
+    def group_weldment(self):
+        grouped_data = {}
+        for key, group in itertools.groupby(sorted(self.boxes_xy_label), self.get_first_letter):
+            grouped_data[key] = list(group)
 
+        # 创建子类实例并添加到大类中
+        for key, group in grouped_data.items():
+            subclass = WeldmentClass(key)
+            for g in group:
+                subclass.addshapelist(g, self.boxes_xy_label.get(g))
 
+            self.add_weldmentclass(subclass)
 
 
 
@@ -166,171 +173,272 @@ class StructureClass:
         self.weldmentclasses.append(weldmentclass)
 
 
+
 # 焊接件类
 class WeldmentClass(StructureClass):
     def __init__(self, name):
-        self.ployclasses = []
+        self.shapelist = []
+        self.xylist = []
+        self.methodclasses = []
         self.name = name
+        self.flaglist = []
+        self.result = None
 
+    def addshapelist(self, shape, box_xy):
+        self.shapelist.append(shape)
+        self.xylist.append(box_xy)
 
-    def add_ploy(self, ployclass):
-        self.ployclasses.append(ployclass)
+    def add_method(self, methodclass):
+        self.methodclasses.append(methodclass)
 
 
-class SSIMClass:
-    def __init__(self, label, x1, y1, x2, y2):
+class SSIMDet:
+    def __init__(self, ref_image, query_image, label, box_xy): # x1, y1, x2, y2
         self.name = 'SSIM'
         self.label = label
-        self.x1 = x1
-        self.x2 = x2
-        self.y1 = y1
-        self.y2 = y2
-        self.result = self.SSIMfunc()
-
-    def SSIMfunc(self):
-        return self.label, self.x1, self.x2, self.y1, self.y2
+        self.x1, self.y1, self.x2, self.y2 = box_xy
+        self.cut_ref_image = self.cut_image(ref_image)
+        self.cut_query_image = self.cut_image(query_image)
+        self.result = self.ssim_func(self.cut_ref_image, self.cut_query_image)
+
+    def cut_image(self, image):
+        return image[self.y1:self.y2, self.x1:self.x2]
+
+    def ssim_func(self, im1, im2):
+
+        imgsize = im1.shape[1] * im1.shape[2]
+        avg1 = im1.mean((1, 2), keepdims=1)
+        avg2 = im2.mean((1, 2), keepdims=1)
+        std1 = im1.std((1, 2), ddof=1)
+        std2 = im2.std((1, 2), ddof=1)
+        cov = ((im1 - avg1) * (im2 - avg2)).mean((1, 2)) * imgsize / (imgsize - 1)
+        avg1 = np.squeeze(avg1)
+        avg2 = np.squeeze(avg2)
+        k1 = 0.01
+        k2 = 0.03
+        c1 = (k1 * 255) ** 2
+        c2 = (k2 * 255) ** 2
+        c3 = c2 / 2
+        # return np.mean((cov + c3) / (std1 * std2  + c3))
+        return np.mean(
+            (2 * avg1 * avg2 + c1) * 2 * (cov + c3) / (avg1 ** 2 + avg2 ** 2 + c1) / (std1 ** 2 + std2 ** 2 + c2))
+
+
+class VarianceDet:
+    def __init__(self, ref_image, query_image, label, box_xy):
+        self.name = 'VarianceDet'
+        self.label = label
+        self.x1, self.y1, self.x2, self.y2 = box_xy
+        self.cut_ref_image = self.cut_image(ref_image)
+        self.cut_query_image = self.cut_image(query_image)
+        self.proportion = self.black_pixels_proportion(self.cut_query_image)
+        if self.proportion > 0.05:
+            self.result = 1
+        else:
+            self.result = self.variance_det_func(self.cut_ref_image, self.cut_query_image)
+
+    def cut_image(self, image):
+        return image[self.y1:self.y2, self.x1:self.x2]
+
+    def black_pixels_proportion(self, cut_query_image):
+        black_pixels = np.sum(cv2.cvtColor(cut_query_image, cv2.COLOR_BGR2GRAY) == 0)
+        other_pixels = np.sum(cv2.cvtColor(cut_query_image, cv2.COLOR_BGR2GRAY) != 0)
+        proportion = black_pixels / (other_pixels + black_pixels)
+        return proportion
+
+    # 计算两张图片的方差
+    def calculate_variance(self, image):
+        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        variance = np.var(gray_image)
+        mean = np.mean(image)
+        variance = variance / mean
+        return variance
+
+    def variance_det_func(self, ref_image, query_image):
+        variance1 = self.calculate_variance(ref_image)
+        variance2 = self.calculate_variance(query_image)
+        variance_diff = abs(variance1 - variance2)
+        max_variance = max(variance1, variance2)
+        normalized_diff = variance_diff / max_variance if max_variance != 0 else 0
+        if normalized_diff < 0.9:
+            return True
+        else:
+            return False
 
 
-class Ploy1Class:
-    def __init__(self, label, x1, y1, x2, y2):
-        self.name = 'Ploy1'
+class RorateDet:
+    def __init__(self, ref_image, query_image, label, box_xy):
+        self.name = 'RorateDet'
         self.label = label
-        self.x1 = x1
-        self.x2 = x2
-        self.y1 = y1
-        self.y2 = y2
-        self.result = self.ploy1func()
+        self.x1, self.y1, self.x2, self.y2 = box_xy
+        self.cut_ref_image = self.cut_image(ref_image)
+        self.cut_query_image = self.cut_image(query_image)
+        self.query_image_rotate = cv2.rotate(self.cut_query_image, cv2.ROTATE_180)
+        self.result = self.rorate_det_func()
 
-    def ploy1func(self):
-        return self.label, self.x1, self.x2, self.y1, self.y2
+    def cut_image(self, image):
+        return image[self.y1:self.y2, self.x1:self.x2]
 
+    # 配准操作，计算模板图与查询图配准点的距离差异（翻转后的距离大于未翻转的）
+    def match_image(self, ref_img, query_img):
+        superglue_matcher = Matcher(
+            {
+                "superpoint": {
+                    "input_shape": (-1, -1),
+                    "keypoint_threshold": 0.003,
+                },
+                "superglue": {
+                    "match_threshold": 0.5,
+                },
+                "use_gpu": True,
+            },shape=20
+        )
+        ref_img = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
+        query_img = cv2.cvtColor(query_img, cv2.COLOR_BGR2GRAY)
+        query_kpts, ref_kpts, _, _, matches = superglue_matcher.match(ref_img, query_img)
+        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)
+        print(diff_query_ref)
+        if len(matches) != 0:
+            diff_query_ref = np.linalg.norm(diff_query_ref, axis=1, ord=2)
+            diff_query_ref = np.sqrt(diff_query_ref)
+            diff_query_ref = np.mean(diff_query_ref)
+        else:
+            diff_query_ref = np.inf
+        return diff_query_ref  # 返回差异值，用于后续做比较
+
+    def rorate_det_func(self):
+        """
+        True: 没有缺陷
+        False：有缺陷
+        :return:
+        """
 
-class Ploy2Class:
-    def __init__(self, label, x1, y1, x2, y2):
-        self.name = 'Ploy2'
-        self.label = label
-        self.x1 = x1
-        self.x2 = x2
-        self.y1 = y1
-        self.y2 = y2
-        self.result = self.ploy2func()
 
-    def ploy2func(self):
-        return self.label, self.x1, self.x2, self.y1, self.y2
+        diff1 = self.match_image(self.cut_ref_image, self.cut_query_image)  # 计算模板图与查询图的配准点差异
+        diff2 = self.match_image(self.cut_ref_image, self.query_image_rotate)  # 计算模板图与翻转180度图的配准点差异
 
+        if diff1 < diff2:
+            return True
+        return False
 
-class Ploy3Class:
-    def __init__(self, label, x1, y1, x2, y2):
-        self.name = 'Ploy3'
+
+class NumPixel:
+    def __init__(self, ref_image, query_image, label, box_xy, threshld=0.15, x_scale = 120, y_scale = 80):
+        self.name = 'NumPixel'
         self.label = label
-        self.x1 = x1
-        self.x2 = x2
-        self.y1 = y1
-        self.y2 = y2
-        self.result = self.ploy3func()
+        self.scale_box = []
+        self.x_scale, self.y_scale, self.threshld = x_scale, y_scale, threshld
+        self.ref_h, self.ref_w, _ = ref_image.shape
 
-    def ploy3func(self):
-        return self.label, self.x1, self.x2, self.y1, self.y2
+        self.x1, self.y1, self.x2, self.y2 = self.big_box(box_xy)
+        # self.cut_ref_image = self.cut_image(ref_image)
+        self.cut_query_image = self.cut_image(query_image)
+        self.ostu_query_image = self.otsu_binarize(self.cut_query_image)
+        self.ostu_query_image = self.ostu_query_image[self.scale_box[1]:self.scale_box[3], self.scale_box[0]:self.scale_box[2]]
+        self.result = self.num_pixel_func(self.ostu_query_image)
 
 
-# 定义一个函数来获取每个元素的首字母
-def get_first_letter(item):
-    return item[0]
 
+    def big_box(self, box_xy):
+        x1, y1, x2, y2 = box_xy
+        nx1, ny1, nx2, ny2 = 0,0,0,0
+        if x1 >= self.x_scale:
+            nx1 = x1-self.x_scale
+            self.scale_box.append(self.x_scale)
+        else:
+            nx1 = 0
+            self.scale_box.append(x1)
 
-if __name__ == '__main__':
-    ref_image_path = './data/yongsheng_image/ref_image/DSC_0452.JPG'
-    query_image_path = './data/yongsheng_image/test_image_query/DSC_0445.JPG'
-    json_path = './data/yongsheng_image/json/DSC_0452.json'
-    save_image_path = './data/yongsheng_image/test_regis_result'
-    json_mask_path = './data/yongsheng_image/json/DSC_0452_mask.json'
-    for filename in os.listdir(image_dir):
+        if y1 >= self.y_scale:
+            ny1 = y1-self.y_scale
+            self.scale_box.append(self.y_scale)
+        else:
+            ny1 = 0
+            self.scale_box.append(y1)
 
+        if x2 + self.x_scale <= self.ref_w:
+            nx2 = x2 + self.x_scale
+            self.scale_box.append(self.scale_box[0]+(x2-x1))
+        else:
+            nx2 = self.ref_w
+            self.scale_box.append(self.scale_box[0]+(x2-x1))
 
+        if y2 + self.y_scale <= self.ref_h:
+            ny2 = y2 + self.y_scale
+            self.scale_box.append(self.scale_box[1]+(y2-y1))
+        else:
+            ny2 = self.ref_h
+            self.scale_box.append(self.scale_box[1]+(y2-y1))
 
+        return nx1, ny1, nx2, ny2
 
 
 
+    def num_pixel_func(self, ostu_query_image):
+        """
+        True: 无缺陷
+        False：有缺陷
+        :return:
+        :rtype:
+        """
 
+        num_pixel_region_query = round((np.sum(ostu_query_image == 0) / (ostu_query_image.shape[0] * ostu_query_image.shape[1])), 2)
+        if num_pixel_region_query >= self.threshld:
+            return True
+        return False
 
-    # struct = StructureClass('./data/yongsheng_image/ref_image/DSC_0452.JPG', './data/yongsheng_image/test_image_query/DSC_0445.JPG', './data/yongsheng_image/json/DSC_0452.json')
-    struct = StructureClass(ref_image_path, query_image_path, json_path, save_image_path, json_mask_path)
 
-    grouped_data = {}
-    for key, group in itertools.groupby(sorted(struct.boxes_xy_label), get_first_letter):
-        grouped_data[key] = list(group)
-
-    # 创建子类实例并添加到大类中
-    for key, group in grouped_data.items():
-        subclass = WeldmentClass(key)
-        for g in group:
-            if len(g) == 1:
-                xy = struct.boxes_xy_label.get(g)
-                ssim = SSIMClass(g, xy[0], xy[1], xy[2], xy[3])
-                subclass.add_ploy(ssim)
-            else:
-                xy = struct.boxes_xy_label.get(g)
-                if str(g).endswith('1'):
-                    poly = Ploy1Class(g, xy[0], xy[1], xy[2], xy[3])
-                elif str(g).endswith('2'):
-                    poly = SSIMClass(g, xy[0], xy[1], xy[2], xy[3])
-                else:
-                    poly = Ploy3Class(g, xy[0], xy[1], xy[2], xy[3])
 
-                subclass.add_ploy(poly)
-        struct.add_weldmentclass(subclass)
+    def cut_image(self, image):
+        return image[self.y1:self.y2, self.x1:self.x2]
 
+    def otsu_binarize(self, image):
+        gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+        gray_image = cv2.equalizeHist(gray_image)
+        # ret1, mask1 = cv2.threshold(gray_image, 0, 255, cv2.THRESH_BINARY +cv2.THRESH_OTSU)
+        ret2, mask = cv2.threshold(gray_image, 60, 255, cv2.THRESH_BINARY)
+        # mask =mask1 & mask2
+        kernel = np.ones((3, 3), np.uint8)
+        mask = cv2.dilate(mask, kernel, iterations=3)
+        mask = cv2.erode(mask, kernel, iterations=3)
 
+        return mask
 
-    w = WeldmentClass('A')
-    struct.add_weldmentclass(w)
+def read_yaml(yaml_file):
+
+    with open(yaml_file, 'r') as file:
+        data = yaml.safe_load(file)
+    return data
+
+def calculate_det(struct, yaml_data):
+    for weldment in struct.weldmentclasses:
+        shapelist = weldment.shapelist
+        xylist = weldment.xylist
+        for i in range(len(shapelist)):
+            for method in yaml_data.get(shapelist[i]):
+                class_obj = globals()[method]
+                instance = class_obj(struct.ref_image, struct.query_image, shapelist[i], xylist[i])
+                weldment.flaglist.append(instance.result)
+                weldment.result = all(weldment.flaglist)
+                weldment.add_method(instance)
 
-    print()
 
 
 
-    # with open('./DSC_0452.json', 'r') as f:
-    #     data = json.load(f)
-    # save_value = {}
-    # for shape in data['shapes']:
-    #     if 'points' in shape:
-    #         shape['points'] = [[int(round(x)), int(round(y))] for x, y in shape['points']]
-    #         x1, y1 = shape['points'][0]
-    #         x2, y2 = shape['points'][1]
-    #         label = shape['label']
-    #         save_value[label] = [x1, y1, x2, y2]
-    #
-    # # 使用groupby函数根据首字母分组
-    # grouped_data = {}
-    # for key, group in itertools.groupby(sorted(save_value), get_first_letter):
-    #     grouped_data[key] = list(group)
-    #
-    # # 打印分组后的结果
-    # for key, group in grouped_data.items():
-    #     print(f"{key}: {group}")
-    #
-    # # 创建大类实例
-    # big_class = StructureClass()
-    # # 创建子类实例并添加到大类中
-    # for key, group in grouped_data.items():
-    #     subclass = WeldmentClass(key)
-    #     for g in group:
-    #         if len(g) == 1:
-    #             xy = save_value.get(g)
-    #             ssim = SSIMClass(g, xy[0], xy[1], xy[2], xy[3])
-    #             subclass.add_ploy(ssim)
-    #         else:
-    #             xy = save_value.get(g)
-    #             r = random.randint(1, 4)
-    #             if r == 1:
-    #                 poly = Ploy1Class(g, xy[0], xy[1], xy[2], xy[3])
-    #             elif r == 2:
-    #                 poly = Ploy2Class(g, xy[0], xy[1], xy[2], xy[3])
-    #             else:
-    #                 poly = Ploy3Class(g, xy[0], xy[1], xy[2], xy[3])
-    #
-    #             subclass.add_ploy(poly)
-    #     big_class.add_weldmentclass(subclass)
-    #
-    # for subclass in big_class.weldmentclasses:
-    #     print(subclass)
+
+
+if __name__ == '__main__':
+    ref_image_path = './data/yongsheng_image/ref_image/image165214-001.jpg'
+    query_image_path = './data/yongsheng_image/test_image_query/image165214-011.jpg'
+    json_path = './data/yongsheng_image/json/image165214-001.json'
+    save_image_path = './data/yongsheng_image/test_regis_result'
+    json_mask_path = './data/yongsheng_image/json/image165214-001_mask.json'
+    # for filename in os.listdir(image_dir):
+
+    struct = StructureClass(ref_image_path, query_image_path, json_path, save_image_path, json_mask_path)
+    yaml_data = read_yaml('./test.yaml')
+    calculate_det(struct, yaml_data.get('image165214-001'))
+
+    print()

test.yaml

@@ -0,0 +1,23 @@
+image165214-001:
+  A-1-0-Va: [VarianceDet]
+  A-2-0-Va: [VarianceDet]
+  A-3-0-Va: [VarianceDet]
+  A-4-0-Va: [VarianceDet]
+  A-4-1-Np: [NumPixel]
+  A-4-2-Np: [NumPixel]
+  A-4-3-Np: [NumPixel]
+  A-4-4-Np: [NumPixel]
+  A-4-5-Np: [NumPixel]
+  A-5-0-VaRo: [VarianceDet, RorateDet]
+  A-6-0-VaRo: [VarianceDet, RorateDet]
+  A-7-0-Va: [VarianceDet]
+  A-8-0-VaRo: [VarianceDet, RorateDet]
+  A-9-0-VaRo: [VarianceDet, RorateDet]
+  A-10-0-Va: [VarianceDet]
+  A-11-0-Va: [VarianceDet]
+  A-12-0-Va: [VarianceDet]
+  A-13-0-VaRo: [VarianceDet, RorateDet]
+  A-14-0-VaRo: [VarianceDet, RorateDet]
+  A-15-0-Va: [VarianceDet]
+  A-16-0-VaRo: [VarianceDet, RorateDet]
+

tests/test.py

@@ -0,0 +1,17 @@
+"""
+# File       : test.py
+# Time       ：2024-06-23 11:59
+# Author     ：FEANGYANG
+# version    ：python 3.7
+# Contact    ：1071082183@qq.com
+# Description：
+"""
+
+import yaml
+
+# 假设你的YAML内容存储在一个名为 data.yaml 的文件中
+with open('../test.yaml', 'r') as file:
+    data = yaml.safe_load(file)
+
+# 打印读取的数据
+print(data)

var_detection.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2024/6/19 0019 下午 12:24
+# @Author  : liudan
+# @File    : var_detection.py
+# @Software: pycharm
+
+
+import os
+
+import cv2
+from PIL import Image
+import numpy as np
+import random
+
+
+# 计算两张图片的方差
+def calculate_variance(image):
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    variance = np.var(gray_image)
+    mean = np.mean(image)
+    variance = variance / mean
+    return variance
+
+
+# 计算两张图的方差差异，ref_image为模板图，query_image为查询图
+def calculate_variance_difference(ref_image, query_image):
+    variance1= calculate_variance(ref_image)
+    variance2 = calculate_variance(query_image)
+    variance_diff = abs(variance1 - variance2)
+    max_variance = max(variance1, variance2)
+    normalized_diff = variance_diff / max_variance if max_variance != 0 else 0
+    # diff = np.abs(image1.astype(float) - image2.astype(float)).var()
+    return normalized_diff
+
+
+
+# 计算Recall等，positive_folder正样本文件夹，negative_folder负样本文件夹，ref_image_dir模板图文件夹
+def detection(positive_folder, negative_folder, ref_image_dir):
+    filename = os.listdir(ref_image_dir)
+    image_files = [f for f in filename if f.endswith(('.jpg', '.jpeg', '.png', '.JPG'))]
+    image_file =random.choice(image_files) # 随机获得一张模板图
+    ref_image_path = os.path.join(ref_image_dir,image_file)
+    # ref_image_path = 'H:\\dataset\\region_samples\\positive_test_img\\A1\\image165246-002_A.jpg'
+    ref_image = cv2.imread(ref_image_path) # 读取模板图
+    print(f"{image_file}\n")
+
+    true_positives = 0
+    false_positives = 0
+    false_negatives = 0
+    true_negatives = 0
+    threshold_score = 0.9
+
+    # 在正样本中计算Recall等
+    positive_false = []
+    for file in os.listdir(positive_folder):
+        query_image_path = os.path.join(positive_folder, file) # 查询图，来自正样本
+        query_image = cv2.imread(query_image_path)
+        variance_score = calculate_variance_difference(ref_image,query_image) # 计算模板图与查询图方差差异
+
+        if variance_score < threshold_score:
+            true_positives += 1
+        else:
+            false_negatives += 1
+            positive_false.append(file)
+
+    print(positive_false)
+
+    # 在负样本中计算Recall等
+    negative_false =[]
+    for file in os.listdir(negative_folder):
+        query_image_path = os.path.join(negative_folder, file) # 查询图，来自负样本
+        query_image = cv2.imread(query_image_path)
+        variance_score = calculate_variance_difference(ref_image,query_image)
+
+        if variance_score > threshold_score:
+            true_negatives += 1
+        else:
+            false_positives += 1
+            negative_false.append((file,variance_score))
+    print(negative_false)
+
+    Accurary = (true_positives + true_negatives) / (true_positives + true_negatives + false_positives + false_negatives)
+    Precision = true_positives / (true_positives + false_positives)
+    Recall = true_positives / (true_positives + false_negatives)
+    F1_score = 2 * (Precision * Recall) / (Precision + Recall)
+
+    print(f"Accurary:{Accurary: .2f}")
+    print(f"Precision: {Precision:.2f}")
+    print(f"Recall: {Recall:.2f}")
+    print(f"F1 Score: {F1_score:.2f}")
+
+if __name__ == "__main__":
+    positive_folder = './data/624_test/positive/Y'
+    negative_folder = './data/624_test/negative/Y'
+    ref_image_dir = './data/624_test/positive/Y'
+
+    detection(positive_folder, negative_folder, ref_image_dir)