587 lines
13 KiB
Python
587 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import hashlib
|
|
import time
|
|
import uuid
|
|
from typing import Dict, Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Utilities
|
|
# ------------------------------------------------------------
|
|
|
|
def load_intrinsics_npz(npz_path: str):
|
|
data = np.load(npz_path)
|
|
|
|
for k in ('camera_matrix', 'mtx', 'K'):
|
|
if k in data:
|
|
K = data[k].astype(np.float32)
|
|
break
|
|
else:
|
|
raise KeyError('Camera matrix not found in npz')
|
|
|
|
for k in ('dist_coeffs', 'dist', 'D'):
|
|
if k in data:
|
|
D = data[k].astype(np.float32).reshape(-1, 1)
|
|
break
|
|
else:
|
|
D = np.zeros((5, 1), dtype=np.float32)
|
|
|
|
return K, D
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def load_robot_vision_config(robot_json_path: str):
|
|
|
|
with open(robot_json_path, 'r', encoding='utf-8') as f:
|
|
robot = json.load(f)
|
|
|
|
vision_config = robot.get('vision_config', {})
|
|
|
|
marker_type = vision_config.get('MarkerType', 'DICT_4X4_250')
|
|
marker_size = float(vision_config.get('MarkerSize', 0.025))
|
|
|
|
return {
|
|
'MarkerType': marker_type,
|
|
'MarkerSize': marker_size
|
|
}
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def get_aruco_detector(dict_name: str):
|
|
|
|
mapping = {
|
|
'DICT_4X4_250': cv2.aruco.DICT_4X4_250,
|
|
'DICT_5X5_100': cv2.aruco.DICT_5X5_100,
|
|
'DICT_6X6_250': cv2.aruco.DICT_6X6_250,
|
|
'DICT_ARUCO_ORIGINAL': cv2.aruco.DICT_ARUCO_ORIGINAL,
|
|
}
|
|
|
|
dict_id = mapping.get(dict_name, cv2.aruco.DICT_4X4_250)
|
|
|
|
dictionary = cv2.aruco.getPredefinedDictionary(dict_id)
|
|
|
|
try:
|
|
params = cv2.aruco.DetectorParameters()
|
|
except Exception:
|
|
params = cv2.aruco.DetectorParameters_create()
|
|
|
|
try:
|
|
detector = cv2.aruco.ArucoDetector(dictionary, params)
|
|
return detector, None
|
|
|
|
except Exception:
|
|
return None, (dictionary, params)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def detect_markers(image, detector_tuple):
|
|
|
|
detector, fallback = detector_tuple
|
|
|
|
if detector is not None:
|
|
|
|
corners, ids, rejected = detector.detectMarkers(image)
|
|
|
|
else:
|
|
|
|
dictionary, params = fallback
|
|
|
|
corners, ids, rejected = cv2.aruco.detectMarkers(
|
|
image,
|
|
dictionary,
|
|
parameters=params
|
|
)
|
|
|
|
return corners, ids, rejected
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def hash_file(path):
|
|
|
|
sha = hashlib.sha256()
|
|
|
|
with open(path, 'rb') as f:
|
|
|
|
while True:
|
|
|
|
chunk = f.read(1024 * 1024)
|
|
|
|
if not chunk:
|
|
break
|
|
|
|
sha.update(chunk)
|
|
|
|
return sha.hexdigest()
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def polygon_mask(shape, polygon):
|
|
|
|
mask = np.zeros(shape, dtype=np.uint8)
|
|
|
|
cv2.fillConvexPoly(
|
|
mask,
|
|
polygon.astype(np.int32),
|
|
255
|
|
)
|
|
|
|
return mask
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def shrink_polygon(points, scale=0.80):
|
|
|
|
center = np.mean(points, axis=0)
|
|
|
|
shrunk = center + (points - center) * scale
|
|
|
|
return shrunk.astype(np.float32)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def compute_sharpness(gray_image, polygon):
|
|
|
|
shrunk = shrink_polygon(polygon, scale=0.80)
|
|
|
|
mask = polygon_mask(gray_image.shape, shrunk)
|
|
|
|
pixels = gray_image[mask == 255]
|
|
|
|
if pixels.size == 0:
|
|
return 0.0
|
|
|
|
temp = np.zeros_like(gray_image)
|
|
temp[mask == 255] = gray_image[mask == 255]
|
|
|
|
lap = cv2.Laplacian(temp, cv2.CV_64F)
|
|
|
|
values = lap[mask == 255]
|
|
|
|
if values.size == 0:
|
|
return 0.0
|
|
|
|
return float(values.var())
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def compute_contrast(gray_image, polygon):
|
|
|
|
shrunk = shrink_polygon(polygon, scale=0.80)
|
|
|
|
mask = polygon_mask(gray_image.shape, shrunk)
|
|
|
|
pixels = gray_image[mask == 255]
|
|
|
|
if pixels.size == 0:
|
|
|
|
return {
|
|
'p05': 0.0,
|
|
'p95': 0.0,
|
|
'dynamic_range': 0.0,
|
|
'mean_gray': 0.0,
|
|
'std_gray': 0.0
|
|
}
|
|
|
|
p05 = float(np.percentile(pixels, 5))
|
|
p95 = float(np.percentile(pixels, 95))
|
|
|
|
return {
|
|
'p05': p05,
|
|
'p95': p95,
|
|
'dynamic_range': float(p95 - p05),
|
|
'mean_gray': float(np.mean(pixels)),
|
|
'std_gray': float(np.std(pixels))
|
|
}
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def compute_edge_ratio(corners):
|
|
|
|
edge_lengths = []
|
|
|
|
for k in range(4):
|
|
|
|
p1 = corners[k]
|
|
p2 = corners[(k + 1) % 4]
|
|
|
|
edge_lengths.append(
|
|
float(np.linalg.norm(p1 - p2))
|
|
)
|
|
|
|
edge_ratio = (
|
|
max(edge_lengths) /
|
|
max(1e-6, min(edge_lengths))
|
|
)
|
|
|
|
return edge_ratio, edge_lengths
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def compute_geometry_metrics(center, corners, width, height):
|
|
|
|
image_center = np.array(
|
|
[width / 2.0, height / 2.0],
|
|
dtype=np.float32
|
|
)
|
|
|
|
dist_center = np.linalg.norm(center - image_center)
|
|
|
|
max_dist = np.linalg.norm(image_center)
|
|
|
|
distance_center_norm = float(
|
|
dist_center / max(1e-6, max_dist)
|
|
)
|
|
|
|
min_x = np.min(corners[:, 0])
|
|
max_x = np.max(corners[:, 0])
|
|
|
|
min_y = np.min(corners[:, 1])
|
|
max_y = np.max(corners[:, 1])
|
|
|
|
border_distance_px = float(min(
|
|
min_x,
|
|
min_y,
|
|
width - max_x,
|
|
height - max_y
|
|
))
|
|
|
|
return {
|
|
'distance_to_center_norm': distance_center_norm,
|
|
'distance_to_border_px': border_distance_px
|
|
}
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def compute_confidence(
|
|
area_px,
|
|
sharpness,
|
|
edge_ratio,
|
|
dynamic_range,
|
|
border_distance_px
|
|
):
|
|
|
|
score = 1.0
|
|
|
|
# area
|
|
score *= min(1.0, area_px / 1500.0)
|
|
|
|
# sharpness
|
|
score *= min(1.0, sharpness / 120.0)
|
|
|
|
# edge distortion
|
|
score *= 1.0 / max(1.0, edge_ratio)
|
|
|
|
# contrast
|
|
score *= min(1.0, dynamic_range / 80.0)
|
|
|
|
# border distance
|
|
score *= min(1.0, max(0.0, border_distance_px) / 50.0)
|
|
|
|
score = max(0.0, min(1.0, score))
|
|
|
|
return float(score)
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
def main():
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
'-i',
|
|
'--image',
|
|
required=True
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-npz',
|
|
'--intrinsics',
|
|
required=True
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-robot',
|
|
'--robot',
|
|
required=True
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-cameraId',
|
|
'--cameraId',
|
|
required=True,
|
|
type=str
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-outDir',
|
|
'--outDir',
|
|
required=True
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(args.outDir, exist_ok=True)
|
|
|
|
# --------------------------------------------------------
|
|
# Load robot vision config
|
|
# --------------------------------------------------------
|
|
|
|
vision_config = load_robot_vision_config(args.robot)
|
|
|
|
marker_type = vision_config['MarkerType']
|
|
marker_size = vision_config['MarkerSize']
|
|
|
|
# --------------------------------------------------------
|
|
# Load image
|
|
# --------------------------------------------------------
|
|
|
|
image = cv2.imread(args.image)
|
|
|
|
if image is None:
|
|
raise RuntimeError(f'Cannot read image: {args.image}')
|
|
|
|
gray = cv2.cvtColor(
|
|
image,
|
|
cv2.COLOR_BGR2GRAY
|
|
)
|
|
|
|
height, width = gray.shape[:2]
|
|
|
|
# --------------------------------------------------------
|
|
# Intrinsics
|
|
# --------------------------------------------------------
|
|
|
|
K, D = load_intrinsics_npz(args.intrinsics)
|
|
|
|
# --------------------------------------------------------
|
|
# Detection
|
|
# --------------------------------------------------------
|
|
|
|
detector_tuple = get_aruco_detector(marker_type)
|
|
|
|
corners_list, ids, rejected = detect_markers(
|
|
gray,
|
|
detector_tuple
|
|
)
|
|
|
|
detections = []
|
|
|
|
# --------------------------------------------------------
|
|
# Valid detections
|
|
# --------------------------------------------------------
|
|
|
|
if ids is not None:
|
|
|
|
ids = ids.flatten().tolist()
|
|
|
|
for i, marker_id in enumerate(ids):
|
|
|
|
corners = corners_list[i].reshape((4, 2)).astype(np.float32)
|
|
|
|
center = corners.mean(axis=0)
|
|
|
|
area_px = float(
|
|
cv2.contourArea(corners)
|
|
)
|
|
|
|
perimeter_px = float(
|
|
cv2.arcLength(corners, True)
|
|
)
|
|
|
|
edge_ratio, edge_lengths = compute_edge_ratio(corners)
|
|
|
|
sharpness = compute_sharpness(
|
|
gray,
|
|
corners
|
|
)
|
|
|
|
contrast = compute_contrast(
|
|
gray,
|
|
corners
|
|
)
|
|
|
|
geometry = compute_geometry_metrics(
|
|
center,
|
|
corners,
|
|
width,
|
|
height
|
|
)
|
|
|
|
confidence = compute_confidence(
|
|
area_px=area_px,
|
|
sharpness=sharpness,
|
|
edge_ratio=edge_ratio,
|
|
dynamic_range=contrast['dynamic_range'],
|
|
border_distance_px=geometry['distance_to_border_px']
|
|
)
|
|
|
|
detection = {
|
|
|
|
'observation_id': str(uuid.uuid4()),
|
|
|
|
'type': 'aruco',
|
|
|
|
'marker_id': int(marker_id),
|
|
|
|
'marker_size_m': marker_size,
|
|
|
|
'image_points_px': corners.tolist(),
|
|
|
|
'center_px': center.tolist(),
|
|
|
|
'quality': {
|
|
|
|
'area_px': area_px,
|
|
|
|
'perimeter_px': perimeter_px,
|
|
|
|
'sharpness': {
|
|
'laplacian_var': sharpness
|
|
},
|
|
|
|
'contrast': contrast,
|
|
|
|
'geometry': geometry,
|
|
|
|
'edge_ratio': edge_ratio,
|
|
|
|
'edge_lengths_px': edge_lengths
|
|
},
|
|
|
|
'confidence': confidence
|
|
}
|
|
|
|
detections.append(detection)
|
|
|
|
# --------------------------------------------------------
|
|
# Rejected candidates
|
|
# --------------------------------------------------------
|
|
|
|
rejected_candidates = []
|
|
|
|
if rejected is not None:
|
|
|
|
for candidate in rejected:
|
|
|
|
pts = candidate.reshape((-1, 2)).astype(np.float32)
|
|
|
|
center = pts.mean(axis=0)
|
|
|
|
area_px = float(
|
|
cv2.contourArea(pts)
|
|
)
|
|
|
|
rejected_candidates.append({
|
|
|
|
'image_points_px': pts.tolist(),
|
|
|
|
'center_px': center.tolist(),
|
|
|
|
'area_px': area_px
|
|
})
|
|
|
|
# --------------------------------------------------------
|
|
# Final output
|
|
# --------------------------------------------------------
|
|
|
|
output = {
|
|
|
|
'schema_version': '1.0',
|
|
|
|
'created_utc': time.strftime(
|
|
'%Y-%m-%dT%H:%M:%SZ',
|
|
time.gmtime()
|
|
),
|
|
|
|
'vision_config': {
|
|
'MarkerType': marker_type,
|
|
'MarkerSize': marker_size
|
|
},
|
|
|
|
'camera': {
|
|
|
|
'camera_id': args.cameraId,
|
|
|
|
'intrinsics_file': os.path.abspath(args.intrinsics),
|
|
|
|
'camera_matrix': K.tolist(),
|
|
|
|
'distortion_coefficients': D.reshape(-1).tolist()
|
|
},
|
|
|
|
'image': {
|
|
|
|
'image_file': os.path.abspath(args.image),
|
|
|
|
'image_sha256': hash_file(args.image),
|
|
|
|
'width_px': int(width),
|
|
|
|
'height_px': int(height)
|
|
},
|
|
|
|
'aruco': {
|
|
|
|
'dictionary': marker_type,
|
|
|
|
'num_detected_markers': len(detections),
|
|
|
|
'num_rejected_candidates': len(rejected_candidates)
|
|
},
|
|
|
|
'detections': detections,
|
|
|
|
'rejected_candidates': rejected_candidates
|
|
}
|
|
|
|
# --------------------------------------------------------
|
|
# Output path
|
|
# --------------------------------------------------------
|
|
|
|
input_filename = os.path.basename(args.image)
|
|
|
|
input_base = os.path.splitext(input_filename)[0]
|
|
|
|
out_json = os.path.join(
|
|
args.outDir,
|
|
f'{input_base}_aruco_detection.json'
|
|
)
|
|
|
|
# --------------------------------------------------------
|
|
# Save JSON
|
|
# --------------------------------------------------------
|
|
|
|
with open(out_json, 'w', encoding='utf-8') as f:
|
|
|
|
json.dump(
|
|
output,
|
|
f,
|
|
indent=2
|
|
)
|
|
|
|
print(f'Saved: {out_json}')
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
|
|
if __name__ == '__main__':
|
|
main() |