Spaces:
Running
Running
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import json | |
| import math | |
| import os | |
| def ransac(image1, image2, detector_type): | |
| """ | |
| Finds the homography matrix using the RANSAC algorithm with the selected feature detector. | |
| """ | |
| gray1 = cv2.cvtColor(image1, cv2.COLOR_RGB2GRAY) | |
| gray2 = cv2.cvtColor(image2, cv2.COLOR_RGB2GRAY) | |
| if detector_type == "SIFT": | |
| detector = cv2.SIFT_create() | |
| matcher = cv2.FlannBasedMatcher(dict(algorithm=1, trees=5), dict(checks=50)) | |
| elif detector_type == "ORB": | |
| detector = cv2.ORB_create() | |
| matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True) | |
| elif detector_type == "BRISK": | |
| detector = cv2.BRISK_create() | |
| matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True) | |
| elif detector_type == "AKAZE": | |
| detector = cv2.AKAZE_create() | |
| matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True) | |
| elif detector_type == "KAZE": | |
| detector = cv2.KAZE_create() | |
| matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True) | |
| else: | |
| return None | |
| kp1, des1 = detector.detectAndCompute(gray1, None) | |
| kp2, des2 = detector.detectAndCompute(gray2, None) | |
| if des1 is None or des2 is None or len(kp1) < 2 or len(kp2) < 2: | |
| return None | |
| try: | |
| if detector_type == "SIFT": | |
| matches = matcher.knnMatch(des1, des2, k=2) | |
| good_matches = [] | |
| if matches: | |
| for m, n in matches: | |
| if m.distance < 0.75 * n.distance: | |
| good_matches.append(m) | |
| else: | |
| matches = matcher.match(des1, des2) | |
| good_matches = sorted(matches, key=lambda x: x.distance) | |
| except cv2.error as e: | |
| print(f"Error during matching: {e}") | |
| return None | |
| if len(good_matches) > 10: | |
| src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2) | |
| dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2) | |
| H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0) | |
| return H | |
| else: | |
| return None | |
| def get_bounding_box_points(json_data): | |
| """ | |
| Extracts and calculates the four corner points of the bounding box, assuming x,y are top-left. | |
| """ | |
| print_area = json_data['printAreas'][0] | |
| x = print_area['position']['x'] | |
| y = print_area['position']['y'] | |
| w = print_area['width'] | |
| h = print_area['height'] | |
| rotation_deg = print_area['rotation'] | |
| points = np.float32([ | |
| [0, 0], | |
| [w, 0], | |
| [w, h], | |
| [0, h] | |
| ]).reshape(-1, 1, 2) | |
| rotation_rad = math.radians(rotation_deg) | |
| cos_theta = math.cos(rotation_rad) | |
| sin_theta = math.sin(rotation_rad) | |
| rotation_matrix = np.array([ | |
| [cos_theta, -sin_theta], | |
| [sin_theta, cos_theta] | |
| ]) | |
| rotated_points = np.dot(points.reshape(-1, 2), rotation_matrix.T) | |
| final_points = rotated_points + np.array([x, y]) | |
| return final_points.reshape(-1, 1, 2) | |
| def process_and_plot_all_detectors(image1_np, image2_np, json_file): | |
| """ | |
| Processes the images with all available detectors and returns image data for display and download. | |
| """ | |
| if image1_np is None or image2_np is None: | |
| return [None] * 6 | |
| try: | |
| with open(json_file.name, 'r') as f: | |
| data = json.load(f) | |
| except Exception as e: | |
| print(f"Error: Could not read JSON file. {e}") | |
| return [None] * 6 | |
| detectors = ["SIFT", "ORB", "BRISK", "AKAZE", "KAZE"] | |
| gallery_images = [] | |
| download_files = [None] * 5 | |
| for i, detector_type in enumerate(detectors): | |
| H = ransac(image1_np, image2_np, detector_type) | |
| if H is not None: | |
| box_points = get_bounding_box_points(data) | |
| output_flat_img = image1_np.copy() | |
| cv2.polylines(output_flat_img, [np.int32(box_points)], isClosed=True, color=(0, 0, 255), thickness=5) | |
| transformed_box_points = cv2.perspectiveTransform(box_points, H) | |
| output_perspective_img = image2_np.copy() | |
| cv2.polylines(output_perspective_img, [np.int32(transformed_box_points)], isClosed=True, color=(0, 0, 255), thickness=5) | |
| fig, axes = plt.subplots(1, 3, figsize=(18, 6)) | |
| axes[0].imshow(cv2.cvtColor(output_flat_img, cv2.COLOR_BGR2RGB)) | |
| axes[0].set_title(f'Original (Flat) - {detector_type}') | |
| axes[0].axis('off') | |
| axes[1].imshow(cv2.cvtColor(image2_np, cv2.COLOR_BGR2RGB)) | |
| axes[1].set_title('Original (Perspective)') | |
| axes[1].axis('off') | |
| axes[2].imshow(cv2.cvtColor(output_perspective_img, cv2.COLOR_BGR2RGB)) | |
| axes[2].set_title('Projected Bounding Box') | |
| axes[2].axis('off') | |
| plt.tight_layout() | |
| file_name = f"result_{detector_type.lower()}.png" | |
| plt.savefig(file_name) | |
| plt.close(fig) | |
| gallery_images.append(file_name) | |
| download_files[i] = file_name | |
| else: | |
| print(f"Warning: Homography matrix could not be found with {detector_type} detector. Skipping this result.") | |
| # We don't append None to the gallery_images list to avoid the error. | |
| # download_files[i] remains None, which is handled correctly by gr.File. | |
| return [gallery_images] + download_files | |
| iface = gr.Interface( | |
| fn=process_and_plot_all_detectors, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Image 1 (Flat)"), | |
| gr.Image(type="numpy", label="Image 2 (Perspective)"), | |
| gr.File(type="filepath", label="JSON File") | |
| ], | |
| outputs=[ | |
| gr.Gallery(label="Results"), | |
| gr.File(label="Download SIFT Result"), | |
| gr.File(label="Download ORB Result"), | |
| gr.File(label="Download BRISK Result"), | |
| gr.File(label="Download AKAZE Result"), | |
| gr.File(label="Download KAZE Result") | |
| ], | |
| title="Homography and Bounding Box Projection with All Detectors", | |
| description="Upload two images and a JSON file to see the bounding box projection for all 5 feature extraction methods. Each result can be downloaded separately." | |
| ) | |
| iface.launch() |