一. 环境准备

1. 安装Anaconda虚拟环境

Anaconda是一个环境容器,也可以叫环境管理器。 作用:可以在Anaconda容器中为python项目创建不同的环境。在各个不同环境中可以安装不同版本的包并且各个环境互不影响。可以在使用不同项目时任意切换所用环境。

  • 下载安装Anaconda
  • 创建虚拟环境

conda create -n env-name python=3.10.0

2. 安装ultralytics组件

Ultralytics是一个开源的目标检测框架,由Ultralytics公司开发和维护。它基于YOLO(You Only Look Once)系列模型,通过引入新功能和改进,进一步提升了目标检测的性能和灵活性。

pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

  • conda安装

conda install -c conda-forge ultralytics

3. 安装pytorch 

PyTorch是一个开源的Python机器学习库,基于Torch库,底层由C++实现,应用于人工智能领域,如计算机视觉和自然语言处理。

PyTorch主要有两大特征:类似于NumPy的张量计算,能在 GPU 或 MPS 等硬件加速器上加速;基于带自动微分系统的深度神经网络。

  • PyTorch官网地址
  • pip安装
  • pytorch需安装2.0.0版本以上,对应的python也应在3.10.0以上

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

  • conda安装

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia 需要对应自己的cuda版本

二. 后端代码

安装好环境之后,就可以使用官网的训练代码对Yolov8进行训练,并得到最后的参数文件best.pt,紧接着就可以编写后端代码进行预测了。

 接下来展示后端代码,后端代码是基于flask框架完成的,非常简单,首先对训练好的模型进行加载,然后对前端传进来的图像信息进行编码判断,使用加载的模型对图像进行预测,将预测得到的边界框信息转化成json格式回传给前端进行展示。

import cv2
import json
from flask import Flask, request
from PIL import Image
import numpy as np
import base64
import io
import os
from ultralytics import YOLO



PATH_TO_CKPT = '模型文件.pt'
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

model = YOLO(PATH_TO_CKPT)

app = Flask(__name__)

@app.route('/api/', methods=["POST"])
def main_interface():
    response = request.get_json()
    imageData = False;

    if 'image' in response:
        imageData = True;
        data_str = response['image']

    if imageData:
        point = data_str.find(',')
        base64_str = data_str[point:]  # remove unused part like this: "data:image/jpeg;base64,"

        image = base64.b64decode(base64_str)
        img = Image.open(io.BytesIO(image))

        if (img.mode != 'RGB'):
            img = img.convert("RGB")

        img_arr = np.array(img)
        results = model.predict(img_arr, save=True, imgsz=640, conf=0.5)

        return results[0].tojson()

@app.after_request
def add_headers(response):
    response.headers.add('Access-Control-Allow-Origin', '*')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
    return response


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0')

三. 前端代码

使用后端代码将物体的边界框坐标转化成json后回传到前端,在前端代码中,对接收到的边界框坐标信息json字符串进行处理,在前端界面上进行展示。

//使用ajax将前端的图像数据传回给后端,并处理后端回传的边界框坐标信息
function communicate(imgBase64) {
  $.ajax({
    url: URL,
    type: "POST",
    contentType: "application/json",
    data: JSON.stringify({"image": img_base64_url}),
    dataType: "json"
  }).done(function(response_data) {
      Result(response_data);
  });
}


//处理前端回传的图像文件,并调用ajax传输给后端处理
function parseFiles(files) {
  const file = files[0];
  const imageType = /image.*/;
  if (file.type.match(imageType)) {
    const reader = new FileReader();
    reader.readAsDataURL(file);
    reader.onloadend = () => {
      image.src = reader.result;
      communicate(reader.result);
    }
  }
}


//在前端界面上将边界框信息标注到图像上
function drawResult(results) {
    canvas.width = image.width;
    canvas.height = image.height;
    ctx = canvas.getContext('2d');
    ctx.clearRect(0, 0, canvas.width, canvas.height);
    ctx.drawImage(image, 0, 0);
    for(bboxInfo of results) {
      bbox = bboxInfo['box'];
      class_name = bboxInfo['name'];
      score = bboxInfo['confidence'];

      ctx.beginPath();
      ctx.lineWidth="4";

      ctx.strokeStyle="red";
      ctx.fillStyle="red";
      ctx.rect(bbox['x1'], bbox['y1'], bbox['x2'] - bbox['x1'], bbox['y2'] - bbox['y1']);
      ctx.stroke();
      
      ctx.font="30px Arial";
      
      let content = class_name + " " + parseFloat(score).toFixed(2);

      ctx.fillText(content, bbox['x1'], bbox['y1'] < 20 ? bbox['y1'] + 30 : bbox['y1']-5);
  }
}

 以上代码仅仅包含js中处理前端数据和后端数据重要的步骤,希望能对大家搭建前后端目标检测系统能有一定的帮助,如果需要全部代码的话可以私信我。

四. 效果展示

以下是这个前后端目标检测系统的一些大致的展示效果。

 

五. 河面垃圾数据集(9000+张)

河面垃圾数据集共包括有9000+张图像数据,其中:

  • 船载视角河面垃圾图像数据有4000+张
  • 河岸垃圾的图像数据有2000+张
  • 无人机视角下的河面垃圾图像数据有3000+张 

船载视角河面垃圾图像数据如下:

 

 

河岸垃圾的图像数据如下:

 

 

 无人机视角下的河面垃圾图像数据如下:

 

 大家如果对数据集有需要的话请私信我!

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐