tf.expand_dims()用法详解

tf.expand_dims(output_pre, axis=-1)

import tensorflow as tf import numpy as np import cv2 import os import json from tqdm import tqdm class ObjectRecognitionDeployer: def __init__(self, model_path, class_labels): """ 初始化部署器 :param model_path: 模型文件路径 (Keras或TFLite) :param class_labels: 类别标签列表 """ self.class_labels = class_labels self.model_path = model_path self.interpreter = None self.input_details = None self.output_details = None # 根据模型类型加载 if model_path.endswith('.tflite'): self.load_tflite_model(model_path) else: self.model = tf.keras.models.load_model(model_path) self.input_shape = self.model.input_shape[1:3] def load_tflite_model(self, model_path): """加载并配置TFLite模型""" # 加载模型 self.interpreter = tf.lite.Interpreter(model_path=model_path) self.interpreter.allocate_tensors() # 获取输入输出详细信息 self.input_details = self.interpreter.get_input_details() self.output_details = self.interpreter.get_output_details() # 保存输入形状 self.input_shape = tuple(self.input_details[0]['shape'][1:3]) # 安全地打印模型元数据 self.print_model_metadata(model_path) def print_model_metadata(self, model_path): """安全地打印TFLite模型元数据""" try: from tflite_support import metadata displayer = metadata.MetadataDisplayer.with_model_file(model_path) print("--- 模型元数据 ---") print(displayer.get_metadata_json()) print("--- 关联文件 ---") print(displayer.get_packed_associated_file_list()) except (ImportError, ValueError) as e: print(f"警告: 无法获取模型元数据 - {str(e)}") print("使用输入/输出详细信息代替:") print(f"输入: {self.input_details}") print(f"输出: {self.output_details}") def preprocess_image(self, image, input_size, input_dtype=np.float32): """ 预处理图像 :param image: 输入图像 (numpy数组或文件路径) :param input_size: 模型输入尺寸 (height, width) :param input_dtype: 期望的输入数据类型 :return: 预处理后的图像张量 """ if isinstance(image, str): if not os.path.exists(image): raise FileNotFoundError(f"图像文件不存在: {image}") img = cv2.imread(image) if img is None: raise ValueError(f"无法读取图像: {image}") else: img = image # 调整尺寸和颜色空间 img = cv2.resize(img, (input_size[1], input_size[0])) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 根据数据类型进行归一化 if input_dtype == np.uint8: img = img.astype(np.uint8) # 量化模型使用uint8 else: # 浮点模型使用float32 img = img.astype(np.float32) / 255.0 # 添加批次维度 img = np.expand_dims(img, axis=0) return img def predict(self, image): """ 执行预测 :param image: 输入图像 (numpy数组或文件路径) :return: 预测结果 (类别名称, 置信度) """ if self.interpreter is not None: # TFLite模型推理 return self.predict_tflite(image) else: # Keras模型推理 return self.predict_keras(image) def predict_keras(self, image): """使用Keras模型预测""" # 预处理 img = self.preprocess_image(image, self.input_shape, np.float32) # 预测 predictions = self.model.predict(img, verbose=0)[0] class_idx = np.argmax(predictions) confidence = predictions[class_idx] class_name = self.class_labels[class_idx] return class_name, confidence def predict_tflite(self, image): """使用TFLite模型预测""" # 获取输入数据类型 input_dtype = self.input_details[0]['dtype'] # 预处理 img = self.preprocess_image(image, self.input_shape, input_dtype) # 设置输入张量 self.interpreter.set_tensor(self.input_details[0]['index'], img) # 执行推理 self.interpreter.invoke() # 获取输出 output_data = self.interpreter.get_tensor(self.output_details[0]['index']) predictions = output_data[0] # 解析结果 class_idx = np.argmax(predictions) confidence = predictions[class_idx] # 如果输出是量化数据，需要反量化 if self.output_details[0]['dtype'] == np.uint8: # 反量化输出 scale, zero_point = self.output_details[0]['quantization'] confidence = scale * (confidence - zero_point) class_name = self.class_labels[class_idx] return class_name, confidence def benchmark(self, image, runs=100): """ 模型性能基准测试 :param image: 测试图像 :param runs: 运行次数 :return: 平均推理时间(ms), 内存占用(MB) """ # 预热运行 self.predict(image) # 计时测试 start_time = tf.timestamp() for _ in range(runs): self.predict(image) end_time = tf.timestamp() avg_time_ms = (end_time - start_time).numpy() * 1000 / runs # 内存占用 if self.interpreter: # 计算输入张量内存占用 input_size = self.input_details[0]['shape'] dtype_size = np.dtype(self.input_details[0]['dtype']).itemsize mem_usage = np.prod(input_size) * dtype_size / (1024 * 1024) else: # 估算Keras模型内存 mem_usage = self.model.count_params() * 4 / (1024 * 1024) # 假设32位浮点数 return avg_time_ms, mem_usage def create_metadata(self, output_path): """ 创建并保存模型元数据文件 :param output_path: 元数据文件输出路径 """ metadata = { "model_type": "tflite" if self.model_path.endswith('.tflite') else "keras", "class_labels": self.class_labels, "input_size": self.input_shape, "input_dtype": str(self.input_details[0]['dtype']) if self.interpreter else "float32", "quantization": None } if self.interpreter and self.input_details[0]['dtype'] == np.uint8: metadata["quantization"] = { "input_scale": float(self.input_details[0]['quantization'][0]), "input_zero_point": int(self.input_details[0]['quantization'][1]), "output_scale": float(self.output_details[0]['quantization'][0]), "output_zero_point": int(self.output_details[0]['quantization'][1]) } with open(output_path, 'w') as f: json.dump(metadata, f, indent=4) return metadata def convert_to_tflite_with_metadata(self, output_path, quantize=False, representative_data_dir=None): """ 将Keras模型转换为TFLite格式并添加元数据 :param output_path: 输出TFLite文件路径 :param quantize: 是否进行量化 :param representative_data_dir: 代表性数据集目录 """ if not self.model_path.endswith(('.keras', '.h5')): raise ValueError("需要Keras模型格式进行转换") # 加载Keras模型 keras_model = tf.keras.models.load_model(self.model_path) # 创建转换器 converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) if quantize: # 量化配置 converter.optimizations = [tf.lite.Optimize.DEFAULT] # 设置代表性数据集生成器 converter.representative_dataset = lambda: self.representative_dataset( representative_data_dir, input_size=self.input_shape ) # 设置输入输出类型 converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.inference_input_type = tf.uint8 converter.inference_output_type = tf.uint8 # 转换模型 tflite_model = converter.convert() # 保存模型 with open(output_path, 'wb') as f: f.write(tflite_model) print(f"TFLite模型已保存到: {output_path}") # 添加元数据 self.add_tflite_metadata(output_path) return output_path def representative_dataset(self, data_dir=None, input_size=(224, 224), num_samples=100): """ 生成代表性数据集用于量化 :param data_dir: 真实数据目录 :param input_size: 输入尺寸 (height, width) :param num_samples: 样本数量 """ # 优先使用真实数据 if data_dir and os.path.exists(data_dir): image_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] # 限制样本数量 image_files = image_files[:min(len(image_files), num_samples)] print(f"使用 {len(image_files)} 张真实图像进行量化校准") for img_path in tqdm(image_files, desc="量化校准"): try: # 读取并预处理图像 img = cv2.imread(img_path) if img is None: continue img = cv2.resize(img, (input_size[1], input_size[0])) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 # 转换为float32并归一化 img = np.expand_dims(img, axis=0) yield [img] except Exception as e: print(f"处理图像 {img_path} 时出错: {str(e)}") else: # 使用随机数据作为备选 print(f"使用随机数据生成 {num_samples} 个样本进行量化校准") for _ in range(num_samples): # 生成随机图像，归一化到[0,1]范围，使用float32类型 data = np.random.rand(1, input_size[0], input_size[1], 3).astype(np.float32) yield [data] def add_tflite_metadata(self, model_path): """为TFLite模型添加元数据""" # 创建标签文件 labels_path = os.path.join(os.path.dirname(model_path), "labels.txt") with open(labels_path, 'w') as f: for label in self.class_labels: f.write(f"{label}\n") # 创建元数据 metadata_path = os.path.join(os.path.dirname(model_path), "metadata.json") self.create_metadata(metadata_path) print(f"元数据已创建: {metadata_path}") print(f"标签文件已创建: {labels_path}") # 使用示例 if __name__ == "__main__": # 类别标签 CLASS_LABELS = ['book', 'cup', 'glasses', 'phone', 'shoe'] # 初始化部署器 deployer = ObjectRecognitionDeployer( model_path='optimized_model.keras', class_labels=CLASS_LABELS ) # 转换为带元数据的TFLite格式 tflite_path = 'model_quantized.tflite' # 使用真实数据目录进行量化校准 REPRESENTATIVE_DATA_DIR = 'path/to/representative_dataset' # 替换为实际路径 deployer.convert_to_tflite_with_metadata( tflite_path, quantize=True, representative_data_dir=REPRESENTATIVE_DATA_DIR ) # 重新加载带元数据的模型 tflite_deployer = ObjectRecognitionDeployer( model_path=tflite_path, class_labels=CLASS_LABELS ) # 测试预测 test_image = 'test_image.jpg' class_name, confidence = tflite_deployer.predict(test_image) print(f"预测结果: {class_name}, 置信度: {confidence:.2f}") # 性能测试 avg_time, mem_usage = tflite_deployer.benchmark(test_image) print(f"平均推理时间: {avg_time:.2f} ms") print(f"内存占用: {mem_usage:.2f} MB") # 创建元数据文件 metadata = deployer.create_metadata('model_metadata.json') print("模型元数据:", json.dumps(metadata, indent=4)) import tensorflow as tf import numpy as np import cv2 import os import json from tqdm import tqdm class ObjectRecognitionDeployer: def __init__(self, model_path, class_labels): """ 初始化部署器 :param model_path: 模型文件路径 (Keras或TFLite) :param class_labels: 类别标签列表 """ self.class_labels = class_labels self.model_path = model_path self.interpreter = None self.input_details = None self.output_details = None # 根据模型类型加载 if model_path.endswith('.tflite'): self.load_tflite_model(model_path) else: self.model = tf.keras.models.load_model(model_path) self.input_shape = self.model.input_shape[1:3] def load_tflite_model(self, model_path): """加载并配置TFLite模型""" # 加载模型 self.interpreter = tf.lite.Interpreter(model_path=model_path) self.interpreter.allocate_tensors() # 获取输入输出详细信息 self.input_details = self.interpreter.get_input_details() self.output_details = self.interpreter.get_output_details() # 保存输入形状 self.input_shape = tuple(self.input_details[0]['shape'][1:3]) # 安全地打印模型元数据 self.print_model_metadata(model_path) def print_model_metadata(self, model_path): """安全地打印TFLite模型元数据""" try: from tflite_support import metadata displayer = metadata.MetadataDisplayer.with_model_file(model_path) print("--- 模型元数据 ---") print(displayer.get_metadata_json()) print("--- 关联文件 ---") print(displayer.get_packed_associated_file_list()) except (ImportError, ValueError) as e: print(f"警告: 无法获取模型元数据 - {str(e)}") print("使用输入/输出详细信息代替:") print(f"输入: {self.input_details}") print(f"输出: {self.output_details}") def preprocess_image(self, image, input_size, input_dtype=np.float32): """ 预处理图像 :param image: 输入图像 (numpy数组或文件路径) :param input_size: 模型输入尺寸 (height, width) :param input_dtype: 期望的输入数据类型 :return: 预处理后的图像张量 """ if isinstance(image, str): if not os.path.exists(image): raise FileNotFoundError(f"图像文件不存在: {image}") img = cv2.imread(image) if img is None: raise ValueError(f"无法读取图像: {image}") else: img = image # 调整尺寸和颜色空间 img = cv2.resize(img, (input_size[1], input_size[0])) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 根据数据类型进行归一化 if input_dtype == np.uint8: img = img.astype(np.uint8) # 量化模型使用uint8 else: # 浮点模型使用float32 img = img.astype(np.float32) / 255.0 # 添加批次维度 img = np.expand_dims(img, axis=0) return img def predict(self, image): """ 执行预测 :param image: 输入图像 (numpy数组或文件路径) :return: 预测结果 (类别名称, 置信度) """ if self.interpreter is not None: # TFLite模型推理 return self.predict_tflite(image) else: # Keras模型推理 return self.predict_keras(image) def predict_keras(self, image): """使用Keras模型预测""" # 预处理 img = self.preprocess_image(image, self.input_shape, np.float32) # 预测 predictions = self.model.predict(img, verbose=0)[0] class_idx = np.argmax(predictions) confidence = predictions[class_idx] class_name = self.class_labels[class_idx] return class_name, confidence def predict_tflite(self, image): """使用TFLite模型预测""" # 获取输入数据类型 input_dtype = self.input_details[0]['dtype'] # 预处理 img = self.preprocess_image(image, self.input_shape, input_dtype) # 设置输入张量 self.interpreter.set_tensor(self.input_details[0]['index'], img) # 执行推理 self.interpreter.invoke() # 获取输出 output_data = self.interpreter.get_tensor(self.output_details[0]['index']) predictions = output_data[0] # 解析结果 class_idx = np.argmax(predictions) confidence = predictions[class_idx] # 如果输出是量化数据，需要反量化 if self.output_details[0]['dtype'] == np.uint8: # 反量化输出 scale, zero_point = self.output_details[0]['quantization'] confidence = scale * (confidence - zero_point) class_name = self.class_labels[class_idx] return class_name, confidence def benchmark(self, image, runs=100): """ 模型性能基准测试 :param image: 测试图像 :param runs: 运行次数 :return: 平均推理时间(ms), 内存占用(MB) """ # 预热运行 self.predict(image) # 计时测试 start_time = tf.timestamp() for _ in range(runs): self.predict(image) end_time = tf.timestamp() avg_time_ms = (end_time - start_time).numpy() * 1000 / runs # 内存占用 if self.interpreter: # 计算输入张量内存占用 input_size = self.input_details[0]['shape'] dtype_size = np.dtype(self.input_details[0]['dtype']).itemsize mem_usage = np.prod(input_size) * dtype_size / (1024 * 1024) else: # 估算Keras模型内存 mem_usage = self.model.count_params() * 4 / (1024 * 1024) # 假设32位浮点数 return avg_time_ms, mem_usage def create_metadata(self, output_path): """ 创建并保存模型元数据文件 :param output_path: 元数据文件输出路径 """ metadata = { "model_type": "tflite" if self.model_path.endswith('.tflite') else "keras", "class_labels": self.class_labels, "input_size": self.input_shape, "input_dtype": str(self.input_details[0]['dtype']) if self.interpreter else "float32", "quantization": None } if self.interpreter and self.input_details[0]['dtype'] == np.uint8: metadata["quantization"] = { "input_scale": float(self.input_details[0]['quantization'][0]), "input_zero_point": int(self.input_details[0]['quantization'][1]), "output_scale": float(self.output_details[0]['quantization'][0]), "output_zero_point": int(self.output_details[0]['quantization'][1]) } with open(output_path, 'w') as f: json.dump(metadata, f, indent=4) return metadata def convert_to_tflite_with_metadata(self, output_path, quantize=False, representative_data_dir=None): """ 将Keras模型转换为TFLite格式并添加元数据 :param output_path: 输出TFLite文件路径 :param quantize: 是否进行量化 :param representative_data_dir: 代表性数据集目录 """ if not self.model_path.endswith(('.keras', '.h5')): raise ValueError("需要Keras模型格式进行转换") # 加载Keras模型 keras_model = tf.keras.models.load_model(self.model_path) # 创建转换器 converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) if quantize: # 量化配置 converter.optimizations = [tf.lite.Optimize.DEFAULT] # 设置代表性数据集生成器 converter.representative_dataset = lambda: self.representative_dataset( representative_data_dir, input_size=self.input_shape ) # 设置输入输出类型 converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.inference_input_type = tf.uint8 converter.inference_output_type = tf.uint8 # 转换模型 tflite_model = converter.convert() # 保存模型 with open(output_path, 'wb') as f: f.write(tflite_model) print(f"TFLite模型已保存到: {output_path}") # 添加元数据 self.add_tflite_metadata(output_path) return output_path def representative_dataset(self, data_dir=None, input_size=(224, 224), num_samples=100): """ 生成代表性数据集用于量化 :param data_dir: 真实数据目录 :param input_size: 输入尺寸 (height, width) :param num_samples: 样本数量 """ # 优先使用真实数据 if data_dir and os.path.exists(data_dir): image_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] # 限制样本数量 image_files = image_files[:min(len(image_files), num_samples)] print(f"使用 {len(image_files)} 张真实图像进行量化校准") for img_path in tqdm(image_files, desc="量化校准"): try: # 读取并预处理图像 img = cv2.imread(img_path) if img is None: continue img = cv2.resize(img, (input_size[1], input_size[0])) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 # 转换为float32并归一化 img = np.expand_dims(img, axis=0) yield [img] except Exception as e: print(f"处理图像 {img_path} 时出错: {str(e)}") else: # 使用随机数据作为备选 print(f"使用随机数据生成 {num_samples} 个样本进行量化校准") for _ in range(num_samples): # 生成随机图像，归一化到[0,1]范围，使用float32类型 data = np.random.rand(1, input_size[0], input_size[1], 3).astype(np.float32) yield [data] def add_tflite_metadata(self, model_path): """为TFLite模型添加元数据""" # 创建标签文件 labels_path = os.path.join(os.path.dirname(model_path), "labels.txt") with open(labels_path, 'w') as f: for label in self.class_labels: f.write(f"{label}\n") # 创建元数据 metadata_path = os.path.join(os.path.dirname(model_path), "metadata.json") self.create_metadata(metadata_path) print(f"元数据已创建: {metadata_path}") print(f"标签文件已创建: {labels_path}") # 使用示例 if __name__ == "__main__": # 类别标签 CLASS_LABELS = ['book', 'cup', 'glasses', 'phone', 'shoe'] # 初始化部署器 deployer = ObjectRecognitionDeployer( model_path='optimized_model.keras', class_labels=CLASS_LABELS ) # 转换为带元数据的TFLite格式 tflite_path = 'model_quantized.tflite' # 使用真实数据目录进行量化校准 REPRESENTATIVE_DATA_DIR = 'path/to/representative_dataset' # 替换为实际路径 deployer.convert_to_tflite_with_metadata( tflite_path, quantize=True, representative_data_dir=REPRESENTATIVE_DATA_DIR ) # 重新加载带元数据的模型 tflite_deployer = ObjectRecognitionDeployer( model_path=tflite_path, class_labels=CLASS_LABELS ) # 测试预测 test_image = 'test_image.jpg' class_name, confidence = tflite_deployer.predict(test_image) print(f"预测结果: {class_name}, 置信度: {confidence:.2f}") # 性能测试 avg_time, mem_usage = tflite_deployer.benchmark(test_image) print(f"平均推理时间: {avg_time:.2f} ms") print(f"内存占用: {mem_usage:.2f} MB") # 创建元数据文件 metadata = deployer.create_metadata('model_metadata.json') print("模型元数据:", json.dumps(metadata, indent=4)) 上述代码我已经成功执行，并且我的ObjectRecognitionDeployer类路径导入代码是from 计算机视觉.test2 import ObjectRecognitionDeployer