当前位置：首页 > news >正文

用TensorFlow训练自己的第一个模型

news 来源：原创 2024/9/20 9:32:53

现在学AI的一个优势就是：前人栽树后人乘凉，很多资料都已完善，而且有很多很棒的开源作品可以学习，感谢大佬们

项目

项目源码地址
视频教程地址
我在大佬的基础上基于此模型还加上了根据特征值缓存进行快速识别的方法，以应对超市某些未能正确识别的场景，针对项目中window.py文件进行修改和补充：

缓存查询方法

 def query_cache(self, image_features):if not cache:return Nonemax_similarity = -1best_label = Nonefor image_id, (cached_features, label) in cache.items():similarity = self.cosine_similarity(image_features, cached_features)if similarity > max_similarity:max_similarity = similaritybest_label = labelif max_similarity >= 0.5:return best_labelelse:return None

余弦相邻计算

 def cosine_similarity(self, features1, features2):dot_product = np.dot(features1.flatten(), features2.flatten())norm_features1 = np.linalg.norm(features1)norm_features2 = np.linalg.norm(features2)return dot_product / (norm_features1 * norm_features2)

缓存更新方法

 def update_cache(self):input_text = self.input_box.text()self.label = input_text or self.labelself.result.setText(self.label)# 如果缓存已满，移除最久未使用的条目if len(cache) >= CACHE_CAPACITY:cache.popitem(last=False)# 添加新条目cache[self.image_id] = (self.image_features, self.label)self.input_box.clear()self.class_names.append(self.label)

获取图片哈希值

    def get_image_id_from_hash(self, img):buffer = img.tobytes()return hashlib.md5(buffer).hexdigest()

预测图片

    def predict_img(self):self.input_box.clear()img = Image.open('images/target.png')  # 读取图片self.image_id = self.get_image_id_from_hash(img)img = np.asarray(img)  # 将图片转化为numpy的数组start_time = time.time()  # 记录开始时间outputs = self.model.predict(img.reshape(1, 224, 224, 3), batch_size=1, )  # 将图片输入模型得到结果end_time = time.time()  # 记录结束时间elapsed_time = end_time - start_time  # 计算时间差print("运行时间：", elapsed_time, "秒")self.image_features = outputsresult = self.query_cache(outputs)self.label = resultif result is None:result_index = int(np.argmax(outputs))result = self.class_names[result_index]  # 获得对应的水果名称self.result.setText(result)self.label = resultelse:self.result.setText(result)  # 在界面上做显示

UI改造

    def initUI(self):main_widget = QWidget()main_layout = QHBoxLayout()font = QFont('楷体', 15)# 主页面，设置组件并在组件放在布局上left_widget = QWidget()left_layout = QVBoxLayout()img_title = QLabel("样本")img_title.setFont(font)img_title.setAlignment(Qt.AlignCenter)self.img_label = QLabel()img_init = cv2.imread(self.to_predict_name)h, w, c = img_init.shapescale = 400 / himg_show = cv2.resize(img_init, (0, 0), fx=scale, fy=scale)cv2.imwrite("images/show.png", img_show)img_init = cv2.resize(img_init, (224, 224))cv2.imwrite('images/target.png', img_init)self.img_label.setPixmap(QPixmap("images/show.png"))left_layout.addWidget(img_title)left_layout.addWidget(self.img_label, 1, Qt.AlignCenter)left_widget.setLayout(left_layout)right_widget = QWidget()right_layout = QVBoxLayout()btn_change = QPushButton(" 上传图片 ")btn_change.clicked.connect(self.change_img)btn_change.setFont(font)btn_predict = QPushButton(" 开始识别 ")btn_predict.setFont(font)btn_predict.clicked.connect(self.predict_img)btn_update = QPushButton(" 更新缓存 ")btn_update.setFont(font)btn_update.clicked.connect(self.update_cache)label_result = QLabel(' 果蔬名称 ')self.result = QLabel("等待识别")label_result.setFont(QFont('楷体', 16))self.result.setFont(QFont('楷体', 24))self.input_box = QLineEdit()self.input_box.setPlaceholderText("请输入内容...")right_layout.addStretch()right_layout.addWidget(label_result, 0, Qt.AlignCenter)right_layout.addStretch()right_layout.addWidget(self.result, 0, Qt.AlignCenter)right_layout.addStretch()right_layout.addWidget(self.input_box)right_layout.addStretch()right_layout.addWidget(btn_change)right_layout.addWidget(btn_predict)right_layout.addWidget(btn_update)right_layout.addStretch()right_widget.setLayout(right_layout)main_layout.addWidget(left_widget)main_layout.addWidget(right_widget)main_widget.setLayout(main_layout)# 关于页面，设置组件并把组件放在布局上label_super = QLabel("作者：cpa")  # todo 更换作者信息label_super.setFont(QFont('楷体', 12))# label_super.setOpenExternalLinks(True)label_super.setAlignment(Qt.AlignRight)# 添加注释self.addTab(main_widget, '主页')self.setTabIcon(0, QIcon('images/主页面.png'))