现有两个训练好的yolov5s的模型,一个是insulator.pt是专门用来识别绝缘子的,另一个是defect.pt是专门用来识别绝缘子上的破损的。我需要搭建一个flask框架,来做一个用于识别绝缘子上的破损的服务器。现在的思路是,用户向我的服务器发送一张图片,我的服务器先用insulator.pt识别出图中绝缘子的位置,将其裁剪出来,然后再用defect.pt识别出上面是否有破损,最后返回给用户yes or no,如果是yes的话还需要标出破损的位置。请给出搭建这个服务器的每一个具体代码,具体步骤,以保证可以让一个没接触过电脑的人也能根据你的步骤完美的搭建出一个这样的服务器。基于python的flask框架,我的电脑上有postman这个软件。
时间: 2024-02-11 07:08:56 浏览: 40
好的,这是一个比较复杂的任务,创建一个用于识别绝缘子上的破损的服务器需要涉及到很多步骤和代码。下面是实现这个任务的步骤:
1. 安装必要的软件和库
在你的电脑上安装 Python3、Flask 和 PyTorch 的 Python API。还需要安装 Pillow 和 requests 库。
2. 下载 YOLOv5 模型
从 GitHub 上下载 YOLOv5 模型。在这个任务中,我们需要下载两个模型:insulator.pt 和 defect.pt。将它们放在一个名为 models 的文件夹中。
3. 编写 Python 脚本
接下来,我们需要编写 Python 脚本来完成以下任务:
- 接受用户上传的图片;
- 使用 insulator.pt 模型识别图像中的绝缘子位置;
- 裁剪出绝缘子图像;
- 使用 defect.pt 模型检查绝缘子图像是否有破损;
- 如果有破损,标出破损的位置;
- 返回结果给用户。
下面是一个示例代码:
```python
import os
from io import BytesIO
import base64
from PIL import Image
import requests
import torch
from flask import Flask, jsonify, request
from yolov5.utils.general import non_max_suppression
from yolov5.models.experimental import attempt_load
from yolov5.utils.torch_utils import select_device, time_synchronized
app = Flask(__name__)
# 加载模型
model_insulator = attempt_load(os.path.join('models', 'insulator.pt'), map_location='cpu')
model_defect = attempt_load(os.path.join('models', 'defect.pt'), map_location='cpu')
device = select_device('')
# 定义类别
class_names = ['insulator']
class_names_defect = ['defect']
# 定义检测函数
def detect(image, model, class_names):
img_size = model.img_size
img = image.copy()
# 图像预处理
img = img.resize((img_size, img_size))
img = torch.from_numpy(img).to(device).float()
img /= 255.0
img = img.unsqueeze(0)
# 检测
pred = model(img, augment=False)[0]
pred = non_max_suppression(pred, 0.4, 0.5, classes=None, agnostic=False)
# 解析结果
results = []
for det in pred:
if det is not None and len(det):
for *xyxy, conf, cls in reversed(det):
xyxy = [int(x) for x in xyxy]
result = {'class': class_names[int(cls)], 'confidence': conf.item(), 'bbox': xyxy}
results.append(result)
return results
# 定义检测函数(破损)
def detect_defect(image, bbox, model, class_names):
img_size = model.img_size
img = image.copy()
# 图像预处理
img = img.crop(bbox)
img = img.resize((img_size, img_size))
img = torch.from_numpy(img).to(device).float()
img /= 255.0
img = img.unsqueeze(0)
# 检测
pred = model(img, augment=False)[0]
pred = non_max_suppression(pred, 0.4, 0.5, classes=None, agnostic=False)
# 解析结果
results = []
for det in pred:
if det is not None and len(det):
for *xyxy, conf, cls in reversed(det):
xyxy = [int(x) for x in xyxy]
xyxy[0] += bbox[0]
xyxy[1] += bbox[1]
xyxy[2] += bbox[0]
xyxy[3] += bbox[1]
result = {'class': class_names[int(cls)], 'confidence': conf.item(), 'bbox': xyxy}
results.append(result)
return results
# 定义 API 路径
@app.route('/predict', methods=['POST'])
def predict():
# 解析请求
data = request.get_json()
image_data = data['image']
# 解码图像
image = Image.open(BytesIO(base64.b64decode(image_data)))
# 检测绝缘子
results = detect(image, model_insulator, class_names)
if len(results) == 0:
return jsonify({'result': 'no insulator detected'})
# 裁剪绝缘子
bbox = results[0]['bbox']
insulator_image = image.crop(bbox)
# 检测破损
results_defect = detect_defect(image, bbox, model_defect, class_names_defect)
if len(results_defect) == 0:
return jsonify({'result': 'no defect detected'})
# 标出破损
img_draw = ImageDraw.Draw(image)
for result in results_defect:
bbox = result['bbox']
img_draw.rectangle(bbox, outline='red', width=5)
# 返回结果
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
return jsonify({'result': 'defect detected', 'image': img_str})
if __name__ == '__main__':
app.run()
```
4. 测试 API
启动 flask 服务,然后使用 Postman 工具向 /predict API 发送 POST 请求,包含一张待检测的图片。API 将返回一个 JSON 对象,包含检测结果和标记后的图片。
相关推荐
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![rar](https://img-home.csdnimg.cn/images/20210720083606.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)