smartpi 0.1.34__py3-none-any.whl → 0.1.36__py3-none-any.whl

smartpi/posenet_utils.py

@@ -0,0 +1,222 @@
+import tensorflow as tf
+import cv2
+import numpy as np
+import os
+import sys
+from pathlib import Path
+
+# 获取当前脚本的绝对路径
+script_dir = os.path.dirname(os.path.abspath(__file__))
+# 获取项目根目录（假设posenet_utils.py位于pose/lib目录下）
+project_root = os.path.abspath(os.path.join(script_dir, '..'))
+
+# 全局变量存储模型解释器和相关信息
+_interpreter = None
+_input_details = None
+_output_details = None
+# 使用绝对路径定义模型路径
+_MODEL_PATH = os.path.join(project_root, 'posemodel', 'posenet.tflite')  # 默认模型路径
+
+# 人体姿态判断参数（可根据需求调整）
+POSE_THRESHOLD = 0.3  # 单个关键点分数阈值
+REQUIRED_KEYPOINTS = 3  # 判断存在人体所需的有效关键点数量
+# 关键人体关节点索引（对应COCO数据集17个关键点）
+KEY_KEYPOINTS = [0, 1, 2, 3, 4, 5, 6, 7]  # 头部、颈部、肩膀、肘部等关键节点
+
+
+def _load_posenet_model(model_path):
+    """内部函数：加载Posenet TFLite模型"""
+    try:
+        interpreter = tf.lite.Interpreter(model_path=model_path)
+        interpreter.allocate_tensors()
+        input_details = interpreter.get_input_details()
+        output_details = interpreter.get_output_details()
+        return interpreter, input_details, output_details
+    except Exception as e:
+        raise FileNotFoundError(f"模型加载失败: {str(e)}")
+
+
+def _preprocess_image(image_path, input_size=(257, 257)):
+    """内部函数：预处理图像，对齐Web端逻辑"""
+    img = cv2.imread(image_path)
+    if img is None:
+        raise FileNotFoundError(f"无法读取图像: {image_path}")
+
+    # 转为RGB
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    return _preprocess_common(img_rgb, input_size)
+
+
+def _preprocess_frame(frame, input_size=(257, 257)):
+    """内部函数：预处理视频帧（numpy数组）"""
+    # 确保输入是BGR格式（OpenCV默认格式）
+    if len(frame.shape) != 3 or frame.shape[2] != 3:
+        raise ValueError(f"无效的帧格式，期望3通道BGR图像，实际为{frame.shape}")
+
+    # 转为RGB
+    img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+    return _preprocess_common(img_rgb, input_size)
+
+
+def _preprocess_common(img_rgb, input_size=(257, 257)):
+    """通用预处理逻辑，供图像和帧处理共享"""
+    # 计算缩放比例
+    scale = min(input_size[0]/img_rgb.shape[1], input_size[1]/img_rgb.shape[0])
+    scaled_width = int(img_rgb.shape[1] * scale)
+    scaled_height = int(img_rgb.shape[0] * scale)
+
+    # 缩放图像（使用线性插值平衡速度和质量）
+    img_scaled = cv2.resize(img_rgb, (scaled_width, scaled_height), interpolation=cv2.INTER_LINEAR)
+
+    # 创建257x257画布，居中放置缩放后的图像
+    img_padded = np.ones((input_size[1], input_size[0], 3), dtype=np.uint8) * 255
+    x_offset = (input_size[0] - scaled_width) // 2
+    y_offset = (input_size[1] - scaled_height) // 2
+    img_padded[y_offset:y_offset+scaled_height, x_offset:x_offset+scaled_width, :] = img_scaled
+
+    # 归一化
+    img_normalized = (img_padded.astype(np.float32) / 127.5) - 1.0
+
+    # 添加批次维度
+    return np.expand_dims(img_normalized, axis=0)
+
+
+def _has_human_pose(heatmap_scores):
+    """判断是否存在人体姿态"""
+    # heatmap_scores形状为 (height, width, num_keypoints)
+    num_keypoints = heatmap_scores.shape[2]
+
+    # 检查关键节点索引是否有效
+    valid_keypoints = [k for k in KEY_KEYPOINTS if k < num_keypoints]
+    if not valid_keypoints:
+        return False, 0
+
+    # 计算每个关键点的最大分数（在整个热图上的最大值）
+    keypoint_max_scores = []
+    for k in valid_keypoints:
+        # 取当前关键点通道的最大分数
+        max_score = np.max(heatmap_scores[..., k])
+        keypoint_max_scores.append(max_score)
+
+    # 统计超过阈值的关键点数量
+    valid_count = sum(1 for score in keypoint_max_scores if score >= POSE_THRESHOLD)
+
+    # 判断是否达到所需数量
+    has_pose = valid_count >= REQUIRED_KEYPOINTS
+    return has_pose, valid_count
+
+
+def get_posenet_output(input_data, model_path=None, output_file=None,
+                       heatmap_file="heatmap.txt", offsets_file="offsets.txt", precision=6):
+    """
+    获取输入的posenet输出，支持图像路径或视频帧(numpy数组)
+
+    参数:
+        input_data: 图像文件路径(str)或视频帧(numpy.ndarray，BGR格式)
+        model_path: 可选，模型文件路径，默认使用 _MODEL_PATH
+        output_file: 可选，拼接后的输出txt文件路径，若为None则不保存
+        heatmap_file: 可选，heatmap数据保存路径，若为None则不保存
+        offsets_file: 可选，offsets数据保存路径，若为None则不保存
+        precision: 数据保存精度（小数位数），默认6位
+
+    返回:
+        元组 (posenet_output, has_pose, valid_keypoint_count)
+            posenet_output: 处理后的一维数组
+            has_pose: 是否检测到人体姿态（bool）
+            valid_keypoint_count: 有效关键点数量
+    """
+    global _interpreter, _input_details, _output_details, _MODEL_PATH
+
+    # 如果指定了新的模型路径或模型未加载，则重新加载模型
+    if model_path is not None or _interpreter is None:
+        model_to_load = model_path if model_path is not None else _MODEL_PATH
+        _interpreter, _input_details, _output_details = _load_posenet_model(model_to_load)
+
+    # 根据输入类型选择预处理方式
+    if isinstance(input_data, str):
+        # 处理图像路径
+        input_tensor = _preprocess_image(input_data)
+    elif isinstance(input_data, np.ndarray):
+        # 处理视频帧（numpy数组）
+        input_tensor = _preprocess_frame(input_data)
+    else:
+        raise TypeError(f"不支持的输入类型: {type(input_data)}，请提供图像路径或numpy数组")
+
+    # 执行推理（复用全局解释器，避免重复初始化）
+    _interpreter.set_tensor(_input_details[0]['index'], input_tensor)
+    _interpreter.invoke()
+
+    # 按名称匹配输出张量
+    output_dict = {}
+    for output in _output_details:
+        output_name = output['name']
+        output_tensor = _interpreter.get_tensor(output['index']).squeeze(axis=0)
+        output_dict[output_name] = output_tensor
+
+    # 提取heatmap和offsets，对heatmap应用Sigmoid激活
+    heatmap = output_dict['MobilenetV1/heatmap_2/BiasAdd']
+    offsets = output_dict['MobilenetV1/offset_2/BiasAdd']
+
+    # 对heatmap应用Sigmoid激活，与TFJS侧的heatmapScores保持一致
+    def sigmoid(x):
+        x = np.clip(x, -500, 500)  # 限制输入范围，防止exp计算溢出
+        return 1 / (1 + np.exp(-x))
+
+    # 生成激活后的heatmap分数（范围[0,1]，与训练数据一致）
+    heatmap_scores = sigmoid(heatmap)
+
+    # 判断是否存在人体姿态
+    has_pose, valid_count = _has_human_pose(heatmap_scores)
+
+    # 拼接激活后的heatmap和offsets（保持与TFJS侧顺序一致）
+    concatenated = np.concatenate([heatmap_scores, offsets], axis=2)
+    posenet_output = concatenated.astype(np.float32).flatten()
+
+    # 保存拼接后的输出（仅当指定了路径且输入是图像时）
+    if output_file is not None and isinstance(input_data, str):
+        output_dir = Path(output_file).parent
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        with open(output_file, 'w', encoding='utf-8') as f:
+            for value in posenet_output:
+                f.write(f"{value:.{precision}f}\n")
+        print(f"拼接后的posenet输出已保存到: {output_file}")
+
+    return posenet_output, has_pose, valid_count
+
+
+# 配套加载函数：从按行保存的txt文件加载数据
+def load_posenet_output(txt_path):
+    """从按行保存的txt文件加载posenet_output"""
+    if not Path(txt_path).exists():
+        raise FileNotFoundError(f"文件不存在: {txt_path}")
+
+    with open(txt_path, 'r', encoding='utf-8') as f:
+        # 读取所有行，跳过空行并转换为float
+        data = [float(line.strip()) for line in f if line.strip() and not line.strip().startswith('shape:')]
+
+    return np.array(data, dtype=np.float32)
+
+
+# 加载heatmap或offsets数据的函数
+def load_posenet_component(txt_path):
+    """从保存的文件加载heatmap或offsets数据，保留原始形状"""
+    if not Path(txt_path).exists():
+        raise FileNotFoundError(f"文件不存在: {txt_path}")
+
+    with open(txt_path, 'r', encoding='utf-8') as f:
+        lines = [line.strip() for line in f if line.strip()]
+
+    # 解析形状信息
+    shape_line = next(line for line in lines if line.startswith('shape:'))
+    shape_str = shape_line.split('shape: ')[1].strip('()')
+    shape = tuple(map(int, shape_str.split(',')))
+
+    # 解析数据
+    data_lines = [line for line in lines if not line.startswith('shape:')]
+    data = np.array([float(line) for line in data_lines], dtype=np.float32)
+
+    # 重塑为原始形状
+    return data.reshape(shape)

smartpi/rknn_hand_workflow.py

@@ -0,0 +1,245 @@
+import cv2
+import numpy as np
+import mediapipe as mp
+import json
+from PIL import Image
+import os
+from rknnlite.api import RKNNLite  # 导入RKNNLite
+import time  # 用于时间测量
+
+class GestureWorkflow:
+    def __init__(self, model_path):
+        # 确保model_path是绝对路径
+        self.model_path = os.path.abspath(model_path)
+        
+        # 初始化MediaPipe Hands
+        self.mp_hands = mp.solutions.hands
+        self.hands = self.mp_hands.Hands(
+            static_image_mode=False,  # 视频流模式  如果只是获取照片的手势关键点 请设置为True
+            max_num_hands=1,#如果想要检测双手，请设置成2
+            min_detection_confidence=0.5,#手势关键点的阈值
+            model_complexity=0#使用最简单的模型  如果效果不准确 可以考虑设置比较复制的模型  1
+        )
+        
+        # 初始化元数据
+        self.min_vals = None
+        self.max_vals = None
+        self.class_labels = None
+        
+        # 加载模型和元数据
+        self.load_model(self.model_path)
+    
+    def load_model(self, model_path):
+        """加载RKNN模型并解析元数据"""
+        # 创建RKNNLite实例
+        self.rknn_lite = RKNNLite()
+        
+        # 加载RKNN模型
+        ret = self.rknn_lite.load_rknn(model_path)
+        if ret != 0:
+            raise RuntimeError(f'加载RKNN模型失败, 错误码: {ret}')
+        
+        # 初始化运行时环境  强制使用npu  core_mask=RKNNLite.NPU_CORE_0
+        ret = self.rknn_lite.init_runtime()
+        if ret != 0:
+            raise RuntimeError(f'初始化NPU运行时失败, 错误码: {ret}')
+        
+        # 从同目录的JSON文件加载元数据
+        metadata_path = self._get_metadata_path(model_path)
+        self._load_metadata(metadata_path)
+    
+    def _get_metadata_path(self, model_path):
+        """获取元数据文件的绝对路径"""
+        # 尝试与模型同目录的JSON文件
+        base_dir = os.path.dirname(model_path)
+        base_name = os.path.basename(model_path)
+        metadata_name = os.path.splitext(base_name)[0] + 'rknn_metadata.json'
+        metadata_path = os.path.join(base_dir, metadata_name)
+        
+        # 如果文件不存在，尝试默认名称
+        if not os.path.exists(metadata_path):
+            metadata_path = os.path.join(base_dir, 'rknn_metadata.json')
+            
+        return metadata_path
+    
+    def _load_metadata(self, metadata_path):
+        """从JSON文件加载元数据"""
+        try:
+            with open(metadata_path, 'r', encoding='utf-8') as f:
+                metadata = json.load(f)
+                
+            self.class_labels = metadata.get('classes', ["点赞", "点踩", "胜利", "拳头", "我爱你", "手掌"])
+            min_max = metadata.get('minMax', {})
+            self.min_vals = min_max.get('min', [])
+            self.max_vals = min_max.get('max', [])
+            
+            print(f"从 {metadata_path} 加载元数据成功")
+            print(f"类别标签: {self.class_labels}")
+            
+        except Exception as e:
+            print(f"加载元数据失败: {e}")
+            # 设置默认值
+            self.class_labels = ["点赞", "点踩", "胜利", "拳头", "我爱你", "手掌"]
+            self.min_vals = []
+            self.max_vals = []
+    
+    def preprocess_image(self, image, target_width=224, target_height=224):
+        """
+        预处理图像：保持比例缩放并居中放置在目标尺寸的画布上
+        返回处理后的OpenCV图像 (BGR格式)
+        """
+        # 将OpenCV图像转换为PIL格式
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(image_rgb)
+        
+        # 计算缩放比例
+        width, height = pil_image.size
+        scale = min(target_width / width, target_height / height)
+        
+        # 计算新尺寸和位置
+        new_width = int(width * scale)
+        new_height = int(height * scale)
+        x = (target_width - new_width) // 2
+        y = (target_height - new_height) // 2
+        
+        # 创建白色背景画布并粘贴缩放后的图像
+        canvas = Image.new('RGB', (target_width, target_height), (255, 255, 255))
+        resized_image = pil_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        canvas.paste(resized_image, (x, y))
+        
+        # 转换回OpenCV格式
+        processed_image = np.array(canvas)
+        return cv2.cvtColor(processed_image, cv2.COLOR_RGB2BGR)
+    
+    def extract_hand_keypoints(self, image):
+        """从图像中提取手部关键点"""
+        # 转换图像为RGB格式并处理
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        results = self.hands.process(image_rgb)
+        
+        if results.multi_hand_landmarks:
+            # 只使用检测到的第一只手
+            landmarks = results.multi_hand_world_landmarks[0]
+            
+            # 提取关键点坐标
+            keypoints = []
+            for landmark in landmarks.landmark:
+                keypoints.extend([landmark.x, landmark.y, landmark.z])
+                
+            return np.array(keypoints, dtype=np.float32)
+        return None
+    
+    def normalize_keypoints(self, keypoints):
+        """归一化关键点数据"""
+        if not self.min_vals or not self.max_vals or len(self.min_vals) != len(keypoints):
+            # 如果没有归一化参数或长度不匹配，返回原始数据
+            return keypoints
+            
+        normalized = []
+        for i, value in enumerate(keypoints):
+            min_val = self.min_vals[i]
+            max_val = self.max_vals[i]
+            if max_val - min_val > 1e-6:  # 避免除以零
+                normalized.append((value - min_val) / (max_val - min_val))
+            else:
+                normalized.append(0.0)
+        
+        return np.array(normalized, dtype=np.float32)
+    
+    def predict_frame(self, frame):
+        """执行手势分类预测（直接处理图像帧）"""
+        # 记录开始时间
+        start_time = time.time()
+        
+        # 预处理图像
+        processed_image = self.preprocess_image(frame, 224, 224)
+        
+        # 提取关键点
+        keypoints = self.extract_hand_keypoints(processed_image)
+        min_time = time.time()
+        hand_time = min_time - start_time      
+        #print(f"关键点识别耗时: {hand_time:.4f}秒")
+        if keypoints is None:
+            # 记录结束时间并计算耗时
+            end_time = time.time()
+            #print(f"识别耗时: {end_time - start_time:.4f}秒 (未检测到手部)")
+            return None, {"error": "未检测到手部", "processing_time": end_time - start_time}
+        
+        # 归一化关键点
+        normalized_kps = self.normalize_keypoints(keypoints)
+        
+        # 准备输入数据 (1, 63) 形状
+        input_data = normalized_kps.reshape(1, -1).astype(np.float32)
+        
+        # 使用RKNN Lite进行推理
+        try:
+            outputs = self.rknn_lite.inference(inputs=[input_data])
+            predictions = outputs[0][0]
+            
+            # 获取预测结果
+            class_id = np.argmax(predictions)
+            confidence = float(predictions[class_id])
+            
+            # 获取类别标签
+            label = self.class_labels[class_id] if class_id < len(self.class_labels) else f"未知类别 {class_id}"
+            
+            # 返回原始结果和格式化结果
+            raw_result = predictions.tolist()
+            formatted_result = {
+                'class': label,
+                'confidence': confidence,
+                'class_id': class_id,
+                'probabilities': raw_result
+            }
+            
+            # 记录结束时间并计算耗时
+            end_time = time.time()
+            rknn_time= end_time - min_time  
+            processing_time  = end_time - start_time   
+            print(f"rknn识别耗时: {rknn_time:.4f}秒") 
+            print(f"总共识别耗时: {processing_time:.4f}秒 - 识别结果: {label} (置信度: {confidence:.2f})")
+            
+            return raw_result, formatted_result
+            
+        except Exception as e:
+            # 记录结束时间并计算耗时
+            end_time = time.time()
+            print(f"推理失败: {e}, 耗时: {end_time - start_time:.4f}秒")
+            return None, {"error": f"推理失败: {str(e)}", "processing_time": end_time - start_time}
+    
+    def release(self):
+        """释放资源"""
+        if hasattr(self, 'rknn_lite'):
+            self.rknn_lite.release()
+            print("NPU资源已释放")
+    
+    def __del__(self):
+        """析构函数自动释放资源"""
+        self.release()
+
+    # 保留原始方法以兼容旧代码
+    def predict(self, image_path):
+        """执行手势分类预测（从文件路径）"""
+        # 确保图像路径是绝对路径
+        absolute_image_path = os.path.abspath(image_path)
+        
+        try:
+            # 使用PIL库读取图像，避免libpng版本问题
+            pil_image = Image.open(absolute_image_path)
+            # 转换为RGB格式
+            rgb_image = pil_image.convert('RGB')
+            # 转换为numpy数组
+            image_array = np.array(rgb_image)
+            # 转换为BGR格式（OpenCV使用的格式）
+            image = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
+            
+            if image is None:
+                raise ValueError(f"无法读取图像: {absolute_image_path}")
+            
+            return self.predict_frame(image)
+        except Exception as e:
+            # 如果PIL失败，尝试使用cv2作为备选
+            image = cv2.imread(absolute_image_path)
+            if image is None:
+                raise ValueError(f"无法读取图像: {absolute_image_path}")
+            return self.predict_frame(image)