更新最新直接用point_diff_main
This commit is contained in:
琉璃月光
@ -1,256 +0,0 @@
|
||||
# detect_fixed.py
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
import platform
|
||||
from rknnlite.api import RKNNLite
|
||||
|
||||
# ====================== 配置区 ======================
|
||||
|
||||
IMAGE_PATH = "11.jpg" # 测试图片
|
||||
MODEL_PATH = "point.rknn"
|
||||
OUTPUT_DIR = "./output_rknn"
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
# 固定参考点 & 缩放因子
|
||||
FIXED_REF_POINT = (535, 605)
|
||||
width_mm, width_px = 70.0, 42
|
||||
height_mm, height_px = 890.0, 507
|
||||
SCALE_X = width_mm / float(width_px)
|
||||
SCALE_Y = height_mm / float(height_px)
|
||||
print(f"[INFO] Scale factors: X={SCALE_X:.3f} mm/px, Y={SCALE_Y:.3f} mm/px")
|
||||
|
||||
IMG_SIZE = (640, 640)
|
||||
|
||||
# 设备树路径（用于自动识别平台）
|
||||
DEVICE_COMPATIBLE_NODE = '/proc/device-tree/compatible'
|
||||
|
||||
|
||||
def get_host():
|
||||
system = platform.system()
|
||||
machine = platform.machine()
|
||||
if system == 'Linux' and machine == 'aarch64':
|
||||
try:
|
||||
with open(DEVICE_COMPATIBLE_NODE, 'r') as f:
|
||||
compatible = f.read()
|
||||
if 'rk3588' in compatible:
|
||||
return 'RK3588'
|
||||
elif 'rk3576' in compatible:
|
||||
return 'RK3576'
|
||||
elif 'rk3562' in compatible:
|
||||
return 'RK3562'
|
||||
else:
|
||||
return 'RK3566_RK3568'
|
||||
except Exception as e:
|
||||
print(f"Read device tree failed: {e}")
|
||||
exit(-1)
|
||||
else:
|
||||
return f"{system}-{machine}"
|
||||
|
||||
|
||||
def letterbox_resize(image, size, bg_color=114):
|
||||
target_w, target_h = size
|
||||
h, w = image.shape[:2]
|
||||
scale = min(target_w / w, target_h / h)
|
||||
new_w, new_h = int(w * scale), int(h * scale)
|
||||
resized = cv2.resize(image, (new_w, new_h))
|
||||
canvas = np.full((target_h, target_w, 3), bg_color, dtype=np.uint8)
|
||||
dx, dy = (target_w - new_w) // 2, (target_h - new_h) // 2
|
||||
canvas[dy:dy + new_h, dx:dx + new_w] = resized
|
||||
return canvas, scale, dx, dy
|
||||
|
||||
|
||||
def safe_sigmoid(x):
|
||||
x = np.clip(x, -50, 50)
|
||||
return 1.0 / (1.0 + np.exp(-x))
|
||||
|
||||
|
||||
def softmax(x):
|
||||
x = x - np.max(x)
|
||||
e = np.exp(x)
|
||||
return e / e.sum()
|
||||
|
||||
|
||||
def dfl_to_xywh(loc, grid_x, grid_y, stride):
|
||||
xywh_ = np.zeros(4)
|
||||
xywh = np.zeros(4)
|
||||
for i in range(4):
|
||||
l = loc[i * 16:(i + 1) * 16]
|
||||
l = softmax(l)
|
||||
xywh_[i] = sum(j * l[j] for j in range(16))
|
||||
|
||||
xywh_[0] = (grid_x + 0.5) - xywh_[0]
|
||||
xywh_[1] = (grid_y + 0.5) - xywh_[1]
|
||||
xywh_[2] = (grid_x + 0.5) + xywh_[2]
|
||||
xywh_[3] = (grid_y + 0.5) + xywh_[3]
|
||||
|
||||
xywh[0] = ((xywh_[0] + xywh_[2]) / 2) * stride
|
||||
xywh[1] = ((xywh_[1] + xywh_[3]) / 2) * stride
|
||||
xywh[2] = (xywh_[2] - xywh_[0]) * stride
|
||||
xywh[3] = (xywh_[3] - xywh_[1]) * stride
|
||||
|
||||
xywh[0] -= xywh[2] / 2
|
||||
xywh[1] -= xywh[3] / 2
|
||||
return xywh
|
||||
|
||||
|
||||
def parse_pose_outputs(outputs, conf_threshold=0.5, dx=0, dy=0, scale=1.0):
|
||||
boxes = []
|
||||
obj_confs = []
|
||||
class_ids = []
|
||||
|
||||
for idx in range(3): # det head
|
||||
det = np.array(outputs[idx])[0]
|
||||
C, H, W = det.shape
|
||||
num_classes = C - 64
|
||||
stride = 640 // H
|
||||
|
||||
for h in range(H):
|
||||
for w in range(W):
|
||||
for c in range(num_classes):
|
||||
conf = safe_sigmoid(det[64 + c, h, w])
|
||||
if conf >= conf_threshold:
|
||||
loc = det[:64, h, w].astype(np.float32)
|
||||
xywh = dfl_to_xywh(loc, w, h, stride)
|
||||
boxes.append(xywh)
|
||||
obj_confs.append(conf)
|
||||
class_ids.append(c)
|
||||
|
||||
if not obj_confs:
|
||||
best_box = np.array([0, 0, 0, 0])
|
||||
class_id = -1
|
||||
obj_conf = 0.0
|
||||
else:
|
||||
max_idx = np.argmax(obj_confs)
|
||||
best_box = boxes[max_idx]
|
||||
class_id = class_ids[max_idx]
|
||||
obj_conf = obj_confs[max_idx]
|
||||
|
||||
x, y, w, h = best_box
|
||||
x = (x - dx) / scale
|
||||
y = (y - dy) / scale
|
||||
w /= scale
|
||||
h /= scale
|
||||
best_box = [x, y, w, h]
|
||||
|
||||
kpt_output = np.array(outputs[3])[0]
|
||||
confs = kpt_output[:, 2, :]
|
||||
best_anchor_idx = np.argmax(np.mean(confs, axis=0))
|
||||
kpt_data = kpt_output[:, :, best_anchor_idx]
|
||||
|
||||
keypoints = []
|
||||
for i in range(kpt_data.shape[0]):
|
||||
x_img, y_img, vis_raw = kpt_data[i]
|
||||
vis_prob = safe_sigmoid(vis_raw)
|
||||
x_orig = (x_img - dx) / scale
|
||||
y_orig = (y_img - dy) / scale
|
||||
keypoints.append([x_orig, y_orig, vis_prob])
|
||||
|
||||
return np.array(keypoints), class_id, obj_conf, best_box
|
||||
|
||||
|
||||
def compute_offset(keypoints, fixed_point, scale_x, scale_y):
|
||||
if len(keypoints) < 2: return None
|
||||
p1, p2 = keypoints[0], keypoints[1]
|
||||
cx = (p1[0] + p2[0]) / 2
|
||||
cy = (p1[1] + p2[1]) / 2
|
||||
dx_mm = (cx - fixed_point[0]) * scale_x
|
||||
dy_mm = (cy - fixed_point[1]) * scale_y
|
||||
return cx, cy, dx_mm, dy_mm
|
||||
|
||||
|
||||
def visualize_result(image, keypoints, bbox, fixed_point, offset_info, save_path):
|
||||
vis = image.copy()
|
||||
fx, fy = map(int, fixed_point)
|
||||
cx, cy, dx_mm, dy_mm = offset_info
|
||||
|
||||
for i, (x, y, conf) in enumerate(keypoints[:2]):
|
||||
if conf > 0.5:
|
||||
color = (0, 0, 255) if i == 0 else (0, 255, 255)
|
||||
cv2.circle(vis, (int(x), int(y)), 6, color, -1)
|
||||
if len(keypoints) >= 2:
|
||||
cv2.line(vis, (int(keypoints[0][0]), int(keypoints[0][1])),
|
||||
(int(keypoints[1][0]), int(keypoints[1][1])),
|
||||
(0, 255, 0), 2)
|
||||
|
||||
x, y, w, h = bbox
|
||||
cv2.rectangle(vis, (int(x), int(y)), (int(x + w), int(y + h)), (255, 0, 0), 2)
|
||||
cv2.circle(vis, (int(cx), int(cy)), 10, (0, 255, 0), 3)
|
||||
cv2.circle(vis, (fx, fy), 12, (255, 255, 0), 3)
|
||||
cv2.arrowedLine(vis, (fx, fy), (int(cx), int(cy)), (255, 255, 0), 2, tipLength=0.05)
|
||||
cv2.putText(vis, f"DeltaX={dx_mm:+.1f}mm", (fx + 30, fy - 30),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
|
||||
cv2.putText(vis, f"DeltaY={dy_mm:+.1f}mm", (fx + 30, fy + 30),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
|
||||
cv2.imwrite(save_path, vis)
|
||||
|
||||
|
||||
def main():
|
||||
host_name = get_host()
|
||||
print(f"[INFO] Running on: {host_name}")
|
||||
|
||||
rknn = RKNNLite()
|
||||
|
||||
# ---- Load Model ----
|
||||
ret = rknn.load_rknn(MODEL_PATH)
|
||||
if ret != 0:
|
||||
print("❌ Failed to load RKNN model!")
|
||||
exit(ret)
|
||||
print("✅ Model loaded successfully.")
|
||||
|
||||
# ---- Init Runtime ----
|
||||
if host_name in ['RK3576', 'RK3588']:
|
||||
ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
||||
else:
|
||||
ret = rknn.init_runtime()
|
||||
if ret != 0:
|
||||
print("❌ Init runtime failed!")
|
||||
rknn.release()
|
||||
exit(ret)
|
||||
print("✅ Runtime initialized.")
|
||||
|
||||
# ---- Preprocess ----
|
||||
ori_img = cv2.imread(IMAGE_PATH)
|
||||
if ori_img is None:
|
||||
print(f"❌ Cannot read image: {IMAGE_PATH}")
|
||||
return
|
||||
|
||||
img_resized, scale, dx, dy = letterbox_resize(ori_img, IMG_SIZE)
|
||||
input_tensor = np.expand_dims(img_resized[..., ::-1], 0).astype(np.uint8) # RGB
|
||||
|
||||
# ---- Inference ----
|
||||
print("🔍 Starting inference...")
|
||||
outputs = rknn.inference(inputs=[input_tensor])
|
||||
print("✅ Inference completed.")
|
||||
|
||||
# ---- Post-process ----
|
||||
try:
|
||||
keypoints, cls_id, obj_conf, bbox = parse_pose_outputs(
|
||||
outputs, dx=dx, dy=dy, scale=scale)
|
||||
|
||||
offset_info = compute_offset(keypoints, FIXED_REF_POINT, SCALE_X, SCALE_Y)
|
||||
if offset_info is None:
|
||||
print("⚠️ Not enough keypoints detected.")
|
||||
return
|
||||
|
||||
cx, cy, dx_mm, dy_mm = offset_info
|
||||
vis_save_path = os.path.join(OUTPUT_DIR, f"result_{os.path.basename(IMAGE_PATH)}")
|
||||
visualize_result(ori_img, keypoints, bbox, FIXED_REF_POINT, offset_info, vis_save_path)
|
||||
|
||||
print(f"\n🎯 Detection Result:")
|
||||
print(f"Center: ({cx:.1f}, {cy:.1f})")
|
||||
print(f"Offset: ΔX={dx_mm:+.2f}mm, ΔY={dy_mm:+.2f}mm")
|
||||
print(f"Class: {cls_id}, Confidence: {obj_conf:.3f}")
|
||||
print(f"Saved result to: {vis_save_path}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Post-processing error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
finally:
|
||||
rknn.release()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,230 +0,0 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
from rknnlite.api import RKNNLite
|
||||
|
||||
# ====================== 配置区 ======================
|
||||
MODEL_PATH = "point.rknn"
|
||||
OUTPUT_DIR = "./output_rknn"
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
# 固定参考点（像素坐标）
|
||||
FIXED_REF_POINT = (535, 605)
|
||||
|
||||
# mm/px 缩放因子（根据标定数据填写）
|
||||
width_mm = 70.0
|
||||
width_px = 42
|
||||
SCALE_X = width_mm / float(width_px)
|
||||
height_mm = 890.0
|
||||
height_px = 507
|
||||
SCALE_Y = height_mm / float(height_px)
|
||||
print(f"Scale factors: SCALE_X={SCALE_X:.3f} mm/px, SCALE_Y={SCALE_Y:.3f} mm/px")
|
||||
|
||||
# 输入尺寸
|
||||
IMG_SIZE = (640, 640)
|
||||
|
||||
|
||||
def letterbox_resize(image, size, bg_color=114):
|
||||
"""保持比例缩放并填充到指定大小"""
|
||||
target_w, target_h = size
|
||||
h, w = image.shape[:2]
|
||||
scale = min(target_w / w, target_h / h)
|
||||
new_w, new_h = int(w * scale), int(h * scale)
|
||||
resized = cv2.resize(image, (new_w, new_h))
|
||||
canvas = np.full((target_h, target_w, 3), bg_color, dtype=np.uint8)
|
||||
dx, dy = (target_w - new_w) // 2, (target_h - new_h) // 2
|
||||
canvas[dy:dy + new_h, dx:dx + new_w] = resized
|
||||
return canvas, scale, dx, dy
|
||||
|
||||
|
||||
def safe_sigmoid(x):
|
||||
x = np.clip(x, -50, 50)
|
||||
return 1.0 / (1.0 + np.exp(-x))
|
||||
|
||||
|
||||
def parse_pose_outputs(outputs, dx=0, dy=0, scale=1.0):
|
||||
"""
|
||||
解析 RKNN YOLO-Pose 关键点输出
|
||||
outputs[3]: shape (1, 4, 3, 8400) -> [kpt_id, (x,y,conf), anchor]
|
||||
"""
|
||||
kpt_output = np.array(outputs[3])[0] # (4, 3, 8400)
|
||||
confs = kpt_output[:, 2, :] # 取每个关键点的 visible_conf
|
||||
mean_conf_per_anchor = np.mean(confs, axis=0) # 每个 anchor 的平均可见性
|
||||
best_anchor_idx = np.argmax(mean_conf_per_anchor)
|
||||
kpt_data = kpt_output[:, :, best_anchor_idx] # (4, 3): x, y, vis_conf
|
||||
|
||||
keypoints = []
|
||||
for i in range(4):
|
||||
x_img = kpt_data[i, 0]
|
||||
y_img = kpt_data[i, 1]
|
||||
vis_conf_raw = kpt_data[i, 2]
|
||||
vis_prob = safe_sigmoid(vis_conf_raw)
|
||||
|
||||
# 映射回原图坐标
|
||||
x_orig = (x_img - dx) / scale
|
||||
y_orig = (y_img - dy) / scale
|
||||
keypoints.append([x_orig, y_orig, vis_prob])
|
||||
|
||||
return np.array(keypoints)
|
||||
|
||||
|
||||
def compute_offset(keypoints, fixed_point, scale_x, scale_y):
|
||||
"""
|
||||
计算中心点相对于固定参考点的偏移量（mm）
|
||||
中心点 = P0 和 P1 的中点
|
||||
返回: (center_x, center_y, dx_mm, dy_mm)
|
||||
"""
|
||||
if len(keypoints) < 2:
|
||||
return None
|
||||
|
||||
p1, p2 = keypoints[0], keypoints[1]
|
||||
cx = (p1[0] + p2[0]) / 2.0
|
||||
cy = (p1[1] + p2[1]) / 2.0
|
||||
|
||||
dx_px = cx - fixed_point[0]
|
||||
dy_px = cy - fixed_point[1]
|
||||
dx_mm = dx_px * scale_x
|
||||
dy_mm = dy_px * scale_y
|
||||
|
||||
return cx, cy, dx_mm, dy_mm
|
||||
|
||||
|
||||
def visualize_result(image, keypoints, fixed_point, offset_info, save_path):
|
||||
"""
|
||||
可视化关键点、参考点、中心点、偏移箭头和文字
|
||||
"""
|
||||
vis = image.copy()
|
||||
colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (255, 255, 0)]
|
||||
cx, cy, dx_mm, dy_mm = offset_info
|
||||
fx, fy = map(int, fixed_point)
|
||||
|
||||
# 绘制关键点
|
||||
for i, (x, y, conf) in enumerate(keypoints):
|
||||
if conf > 0.5:
|
||||
cv2.circle(vis, (int(x), int(y)), 8, colors[i], -1)
|
||||
cv2.putText(vis, f"P{i}", (int(x) + 10, int(y) - 10),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, colors[i], 2)
|
||||
|
||||
# 绘制中心点
|
||||
cv2.circle(vis, (int(cx), int(cy)), 12, (0, 255, 0), 3)
|
||||
cv2.putText(vis, "Center", (int(cx) + 20, int(cy)),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 3)
|
||||
|
||||
# 绘制参考点
|
||||
cv2.circle(vis, (fx, fy), 15, (255, 255, 0), 3)
|
||||
cv2.putText(vis, "Ref", (fx + 20, fy),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 3)
|
||||
|
||||
# 绘制偏移箭头和文字
|
||||
cv2.arrowedLine(vis, (fx, fy), (int(cx), int(cy)), (0, 255, 255), 3, tipLength=0.05)
|
||||
cv2.putText(vis, f"DeltaX={dx_mm:+.1f}mm", (fx + 40, fy - 40),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 3)
|
||||
cv2.putText(vis, f"DeltaY={dy_mm:+.1f}mm", (fx + 40, fy + 40),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 3)
|
||||
|
||||
cv2.imwrite(save_path, vis)
|
||||
|
||||
|
||||
def calculate_offset_from_image(image_path, visualize=False):
|
||||
"""
|
||||
主函数：输入图片路径，输出偏移量 (dx_mm, dy_mm)
|
||||
|
||||
参数:
|
||||
image_path (str): 输入图像路径
|
||||
visualize (bool): 是否保存可视化结果
|
||||
|
||||
返回:
|
||||
dict: {
|
||||
'success': bool,
|
||||
'dx_mm': float or None,
|
||||
'dy_mm': float or None,
|
||||
'cx': float or None, # 中心点 x
|
||||
'cy': float or None, # 中心点 y
|
||||
'message': str
|
||||
}
|
||||
"""
|
||||
# 读取图像
|
||||
orig = cv2.imread(image_path)
|
||||
if orig is None:
|
||||
return {
|
||||
'success': False,
|
||||
'dx_mm': None, 'dy_mm': None, 'cx': None, 'cy': None,
|
||||
'message': f'Failed to load image: {image_path}'
|
||||
}
|
||||
|
||||
h0, w0 = orig.shape[:2]
|
||||
|
||||
# 预处理
|
||||
img_resized, scale, dx, dy = letterbox_resize(orig, IMG_SIZE)
|
||||
infer_img = np.expand_dims(img_resized[..., ::-1], 0).astype(np.uint8)
|
||||
|
||||
# 加载模型并推理
|
||||
rknn = RKNNLite(verbose=False)
|
||||
ret = rknn.load_rknn(MODEL_PATH)
|
||||
if ret != 0:
|
||||
return {
|
||||
'success': False,
|
||||
'dx_mm': None, 'dy_mm': None, 'cx': None, 'cy': None,
|
||||
'message': 'Failed to load RKNN model'
|
||||
}
|
||||
|
||||
try:
|
||||
rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
||||
outputs = rknn.inference([infer_img])
|
||||
except Exception as e:
|
||||
rknn.release()
|
||||
return {
|
||||
'success': False,
|
||||
'dx_mm': None, 'dy_mm': None, 'cx': None, 'cy': None,
|
||||
'message': f'Inference error: {str(e)}'
|
||||
}
|
||||
finally:
|
||||
rknn.release()
|
||||
|
||||
# 解析关键点
|
||||
try:
|
||||
keypoints = parse_pose_outputs(outputs, dx=dx, dy=dy, scale=scale)
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'dx_mm': None, 'dy_mm': None, 'cx': None, 'cy': None,
|
||||
'message': f'Parse keypoint error: {str(e)}'
|
||||
}
|
||||
|
||||
# 计算偏移
|
||||
offset_info = compute_offset(keypoints, FIXED_REF_POINT, SCALE_X, SCALE_Y)
|
||||
if offset_info is None:
|
||||
return {
|
||||
'success': False,
|
||||
'dx_mm': None, 'dy_mm': None, 'cx': None, 'cy': None,
|
||||
'message': 'Not enough keypoints to compute offset'
|
||||
}
|
||||
|
||||
cx, cy, dx_mm, dy_mm = offset_info
|
||||
|
||||
# 可视化（可选）
|
||||
if visualize:
|
||||
vis_save_path = os.path.join(OUTPUT_DIR, f"result_{os.path.basename(image_path)}")
|
||||
visualize_result(orig, keypoints, FIXED_REF_POINT, offset_info, vis_save_path)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'dx_mm': dx_mm,
|
||||
'dy_mm': dy_mm,
|
||||
'cx': cx,
|
||||
'cy': cy,
|
||||
'message': 'Success'
|
||||
}
|
||||
|
||||
|
||||
# ====================== 使用示例 ======================
|
||||
if __name__ == "__main__":
|
||||
image_path = "11.jpg"
|
||||
|
||||
result = calculate_offset_from_image(image_path, visualize=True)
|
||||
|
||||
if result['success']:
|
||||
print(f"Center point: ({result['cx']:.1f}, {result['cy']:.1f})")
|
||||
print(f"Offset: DeltaX={result['dx_mm']:+.2f} mm, DeltaY={result['dy_mm']:+.2f} mm")
|
||||
else:
|
||||
print("Error:", result['message'])
|
||||
@ -1,236 +0,0 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
from rknnlite.api import RKNNLite
|
||||
|
||||
# ====================== 配置区 ======================
|
||||
|
||||
MODEL_PATH = "point.rknn"
|
||||
OUTPUT_DIR = "./output_rknn"
|
||||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||||
|
||||
# 固定参考点（像素坐标）
|
||||
FIXED_REF_POINT = (535, 605)
|
||||
|
||||
# mm/px 缩放因子（根据标定数据填写）
|
||||
width_mm = 70.0
|
||||
width_px = 42
|
||||
SCALE_X = width_mm / float(width_px)
|
||||
height_mm = 890.0
|
||||
height_px = 507
|
||||
SCALE_Y = height_mm / float(height_px)
|
||||
print(f"Scale factors: SCALE_X={SCALE_X:.3f} mm/px, SCALE_Y={SCALE_Y:.3f} mm/px")
|
||||
|
||||
# 输入尺寸
|
||||
IMG_SIZE = (640, 640)
|
||||
|
||||
|
||||
def letterbox_resize(image, size, bg_color=114):
|
||||
target_w, target_h = size
|
||||
h, w = image.shape[:2]
|
||||
scale = min(target_w / w, target_h / h)
|
||||
new_w, new_h = int(w * scale), int(h * scale)
|
||||
resized = cv2.resize(image, (new_w, new_h))
|
||||
canvas = np.full((target_h, target_w, 3), bg_color, dtype=np.uint8)
|
||||
dx, dy = (target_w - new_w) // 2, (target_h - new_h) // 2
|
||||
canvas[dy:dy + new_h, dx:dx + new_w] = resized
|
||||
return canvas, scale, dx, dy
|
||||
|
||||
|
||||
def safe_sigmoid(x):
|
||||
x = np.clip(x, -50, 50)
|
||||
return 1.0 / (1.0 + np.exp(-x))
|
||||
|
||||
|
||||
def softmax(x):
|
||||
x = x - np.max(x)
|
||||
e = np.exp(x)
|
||||
return e / e.sum()
|
||||
|
||||
|
||||
def dfl_to_xywh(loc, grid_x, grid_y, stride):
|
||||
"""将 DFL 输出解析为 xywh"""
|
||||
xywh_ = np.zeros(4)
|
||||
xywh = np.zeros(4)
|
||||
|
||||
# 每个维度 16 bins 做 softmax
|
||||
for i in range(4):
|
||||
l = loc[i * 16:(i + 1) * 16]
|
||||
l = softmax(l)
|
||||
xywh_[i] = sum([j * l[j] for j in range(16)])
|
||||
|
||||
# 对应公式
|
||||
xywh_[0] = (grid_x + 0.5) - xywh_[0]
|
||||
xywh_[1] = (grid_y + 0.5) - xywh_[1]
|
||||
xywh_[2] = (grid_x + 0.5) + xywh_[2]
|
||||
xywh_[3] = (grid_y + 0.5) + xywh_[3]
|
||||
|
||||
# 转成中心点 + 宽高
|
||||
xywh[0] = ((xywh_[0] + xywh_[2]) / 2) * stride
|
||||
xywh[1] = ((xywh_[1] + xywh_[3]) / 2) * stride
|
||||
xywh[2] = (xywh_[2] - xywh_[0]) * stride
|
||||
xywh[3] = (xywh_[3] - xywh_[1]) * stride
|
||||
|
||||
# 转为左上角坐标
|
||||
xywh[0] = xywh[0] - xywh[2] / 2
|
||||
xywh[1] = xywh[1] - xywh[3] / 2
|
||||
return xywh
|
||||
|
||||
|
||||
def parse_pose_outputs(outputs, conf_threshold=0.5, dx=0, dy=0, scale=1.0):
|
||||
"""
|
||||
完整解析 RKNN YOLO-Pose 输出
|
||||
返回 keypoints, class_id, obj_conf, bbox（已映射回原图）
|
||||
"""
|
||||
boxes = []
|
||||
obj_confs = []
|
||||
class_ids = []
|
||||
|
||||
# 遍历前三个输出 tensor (det 输出)
|
||||
for idx in range(3):
|
||||
det = np.array(outputs[idx])[0] # (C,H,W)
|
||||
C, H, W = det.shape
|
||||
num_classes = C - 64 # 前64通道为 DFL bbox
|
||||
stride = 640 // H
|
||||
|
||||
for h in range(H):
|
||||
for w in range(W):
|
||||
for c in range(num_classes):
|
||||
conf = safe_sigmoid(det[64 + c, h, w])
|
||||
if conf >= conf_threshold:
|
||||
loc = det[:64, h, w].astype(np.float32)
|
||||
xywh = dfl_to_xywh(loc, w, h, stride)
|
||||
boxes.append(xywh)
|
||||
obj_confs.append(conf)
|
||||
class_ids.append(c)
|
||||
|
||||
if not obj_confs:
|
||||
best_box = np.array([0, 0, 0, 0])
|
||||
class_id = -1
|
||||
obj_conf = 0.0
|
||||
else:
|
||||
max_idx = np.argmax(obj_confs)
|
||||
best_box = boxes[max_idx]
|
||||
class_id = class_ids[max_idx]
|
||||
obj_conf = obj_confs[max_idx]
|
||||
|
||||
# 🔹 bbox 坐标映射回原图
|
||||
x, y, w, h = best_box
|
||||
x = (x - dx) / scale
|
||||
y = (y - dy) / scale
|
||||
w = w / scale
|
||||
h = h / scale
|
||||
best_box = np.array([x, y, w, h])
|
||||
|
||||
# 🔹 关键点解析
|
||||
kpt_output = np.array(outputs[3])[0] # (num_kpts, 3, num_anchor)
|
||||
confs = kpt_output[:, 2, :]
|
||||
best_anchor_idx = np.argmax(np.mean(confs, axis=0))
|
||||
kpt_data = kpt_output[:, :, best_anchor_idx]
|
||||
|
||||
keypoints = []
|
||||
for i in range(kpt_data.shape[0]):
|
||||
x_img, y_img, vis_conf_raw = kpt_data[i]
|
||||
vis_prob = safe_sigmoid(vis_conf_raw)
|
||||
x_orig = (x_img - dx) / scale
|
||||
y_orig = (y_img - dy) / scale
|
||||
keypoints.append([x_orig, y_orig, vis_prob])
|
||||
|
||||
return np.array(keypoints), class_id, obj_conf, best_box
|
||||
|
||||
|
||||
def compute_offset(keypoints, fixed_point, scale_x, scale_y):
|
||||
if len(keypoints) < 2:
|
||||
return None
|
||||
p1, p2 = keypoints[0], keypoints[1]
|
||||
cx = (p1[0] + p2[0]) / 2.0
|
||||
cy = (p1[1] + p2[1]) / 2.0
|
||||
dx_mm = (cx - fixed_point[0]) * scale_x
|
||||
dy_mm = (cy - fixed_point[1]) * scale_y
|
||||
return cx, cy, dx_mm, dy_mm
|
||||
|
||||
|
||||
def visualize_result(image, keypoints, bbox, fixed_point, offset_info, save_path):
|
||||
vis = image.copy()
|
||||
colors = [(0, 0, 255), (0, 255, 255)]
|
||||
cx, cy, dx_mm, dy_mm = offset_info
|
||||
fx, fy = map(int, fixed_point)
|
||||
|
||||
# 绘制关键点
|
||||
for i, (x, y, conf) in enumerate(keypoints[:2]):
|
||||
if conf > 0.5:
|
||||
cv2.circle(vis, (int(x), int(y)), 6, colors[i], -1)
|
||||
if len(keypoints) >= 2:
|
||||
cv2.line(vis,
|
||||
(int(keypoints[0][0]), int(keypoints[0][1])),
|
||||
(int(keypoints[1][0]), int(keypoints[1][1])),
|
||||
(0, 255, 0), 2)
|
||||
|
||||
# 绘制 bbox
|
||||
x, y, w, h = bbox
|
||||
cv2.rectangle(vis, (int(x), int(y)), (int(x + w), int(y + h)), (255, 0, 0), 2)
|
||||
|
||||
# 绘制中心点
|
||||
cv2.circle(vis, (int(cx), int(cy)), 10, (0, 255, 0), 3)
|
||||
cv2.circle(vis, (fx, fy), 12, (255, 255, 0), 3)
|
||||
cv2.arrowedLine(vis, (fx, fy), (int(cx), int(cy)), (255, 255, 0), 2, tipLength=0.05)
|
||||
cv2.putText(vis, f"DeltaX={dx_mm:+.1f}mm", (fx + 30, fy - 30),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
|
||||
cv2.putText(vis, f"DeltaY={dy_mm:+.1f}mm", (fx + 30, fy + 30),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
|
||||
|
||||
cv2.imwrite(save_path, vis)
|
||||
|
||||
|
||||
def calculate_offset_from_image(image_path, visualize=False):
|
||||
orig = cv2.imread(image_path)
|
||||
if orig is None:
|
||||
return {'success': False, 'message': f'Failed to load image: {image_path}'}
|
||||
|
||||
img_resized, scale, dx, dy = letterbox_resize(orig, IMG_SIZE)
|
||||
infer_img = np.expand_dims(img_resized[..., ::-1], 0).astype(np.uint8)
|
||||
|
||||
rknn = RKNNLite(verbose=False)
|
||||
ret = rknn.load_rknn(MODEL_PATH)
|
||||
if ret != 0:
|
||||
return {'success': False, 'message': 'Failed to load RKNN model'}
|
||||
|
||||
try:
|
||||
rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
||||
outputs = rknn.inference([infer_img])
|
||||
finally:
|
||||
rknn.release()
|
||||
|
||||
try:
|
||||
keypoints, class_id, obj_conf, bbox = parse_pose_outputs(outputs, dx=dx, dy=dy, scale=scale)
|
||||
except Exception as e:
|
||||
return {'success': False, 'message': f'Parse error: {str(e)}'}
|
||||
|
||||
offset_info = compute_offset(keypoints, FIXED_REF_POINT, SCALE_X, SCALE_Y)
|
||||
if offset_info is None:
|
||||
return {'success': False, 'message': 'Not enough keypoints'}
|
||||
|
||||
cx, cy, dx_mm, dy_mm = offset_info
|
||||
|
||||
if visualize:
|
||||
vis_save_path = os.path.join(OUTPUT_DIR, f"result_{os.path.basename(image_path)}")
|
||||
visualize_result(orig, keypoints, bbox, FIXED_REF_POINT, offset_info, vis_save_path)
|
||||
|
||||
return {'success': True, 'dx_mm': dx_mm, 'dy_mm': dy_mm,
|
||||
'cx': cx, 'cy': cy, 'class_id': class_id,
|
||||
'obj_conf': obj_conf, 'bbox': bbox,
|
||||
'message': 'Success'}
|
||||
|
||||
|
||||
# ====================== 使用示例 ======================
|
||||
if __name__ == "__main__":
|
||||
image_path = "11.jpg"
|
||||
result = calculate_offset_from_image(image_path, visualize=True)
|
||||
|
||||
if result['success']:
|
||||
print(f"Center point: ({result['cx']:.1f}, {result['cy']:.1f})")
|
||||
print(f"Offset: DeltaX={result['dx_mm']:+.2f} mm, DeltaY={result['dy_mm']:+.2f} mm")
|
||||
print(f"Class ID: {result['class_id']}, Confidence: {result['obj_conf']:.3f}")
|
||||
print(f"BBox: {result['bbox']}")
|
||||
else:
|
||||
print("Error:", result['message'])
|
||||
@ -37,7 +37,7 @@ def init_rknn_once(model_path):
|
||||
print(f"[ERROR] Failed to load RKNN: {ret}")
|
||||
_rknn_instance = None
|
||||
return None
|
||||
ret = _rknn_instance.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
|
||||
ret = _rknn_instance.init_runtime()
|
||||
if ret != 0:
|
||||
print(f"[ERROR] Failed to init RKNN runtime: {ret}")
|
||||
_rknn_instance = None
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Reference in New Issue
Block a user