如何让Django容器调用Docker化的YoloV5深度学习容器？

一、给YoloV5容器添加API服务 #

YoloV5本身无对外调用接口，需要套一层轻量Web API，推荐用FastAPI（性能适配机器学习场景），步骤如下：

1. 在YoloV5项目根目录创建API服务文件main.py：

from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io
import torch

# 加载预训练YoloV5模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
app = FastAPI()

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # 读取上传图片
    image_bytes = await file.read()
    img = Image.open(io.BytesIO(image_bytes))
    
    # 执行预测
    results = model(img)
    
    # 格式化预测结果（提取类别、置信度、坐标）
    predictions = []
    for pred in results.pandas().xyxy[0].to_dict(orient="records"):
        predictions.append({
            "class": pred["name"],
            "confidence": float(pred["confidence"]),
            "bbox": [
                float(pred["xmin"]),
                float(pred["ymin"]),
                float(pred["xmax"]),
                float(pred["ymax"])
            ]
        })
    
    return {"predictions": predictions}

2. 修改YoloV5的Dockerfile（或新建），安装依赖并启动API：

FROM ultralytics/yolov5:latest

WORKDIR /usr/src/app

# 安装FastAPI及依赖
RUN pip install fastapi uvicorn python-multipart pillow

# 复制API文件
COPY main.py .

# 启动API服务，监听端口8000
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

二、配置Docker自定义网络 #

容器间通信需创建自定义网络，避免默认桥接网络的限制：

1. 创建自定义网络：

docker network create django-yolo-network

所有容器加入该网络后，可通过容器名称直接访问（比如Django容器用yolov5-service作为主机名调用YoloV5 API）。

三、Django容器中调用YoloV5 API #

在Django视图/业务逻辑中，用requests库发送HTTP请求到YoloV5 API，处理返回结果：

1. 在Django项目的requirements.txt中添加requests，确保容器安装该依赖。

2. 编写Django视图示例：

from django.http import JsonResponse
from django.views.decorators.http import require_POST
import requests
import os

@require_POST
def image_predict(request):
    if 'image' not in request.FILES:
        return JsonResponse({"error": "No image uploaded"}, status=400)
    
    image_file = request.FILES['image']
    # 从环境变量读取YoloV5 API地址，避免硬编码
    YOLO_API_URL = os.environ.get("YOLO_API_URL", "http://yolov5-service:8000/predict")
    
    try:
        # 发送图片到YoloV5 API
        response = requests.post(
            YOLO_API_URL,
            files={"file": (image_file.name, image_file, image_file.content_type)}
        )
        response.raise_for_status()
        predictions = response.json()
        
        # 可将预测结果存入Postgres数据库
        # 示例：Prediction.objects.create(image=image_file, results=predictions)
        
        return JsonResponse(predictions)
    except requests.exceptions.RequestException as e:
        return JsonResponse({"error": str(e)}, status=500)

四、用Docker Compose统一管理容器 #

生产环境用Docker Compose简化部署，创建docker-compose.yml：

version: '3.8'

networks:
  django-yolo-network:
    external: true  # 使用之前创建的自定义网络，也可改为internal: true让Compose自动创建

services:
  postgres:
    image: postgres:15-alpine
    environment:
      POSTGRES_DB: django_db
      POSTGRES_USER: django_user
      POSTGRES_PASSWORD: django_pass
    volumes:
      - postgres_data:/var/lib/postgresql/data/
    networks:
      - django-yolo-network
    ports:
      - "5432:5432"

  django:
    build: ./django-app  # 指向你的Django项目目录
    command: gunicorn myproject.wsgi:application --bind 0.0.0.0:8000
    environment:
      - DATABASE_URL=postgres://django_user:django_pass@postgres:5432/django_db
      - YOLO_API_URL=http://yolov5-service:8000/predict
    volumes:
      - ./django-app:/usr/src/app
    networks:
      - django-yolo-network
    ports:
      - "80:8000"
    depends_on:
      - postgres
      - yolov5-service

  yolov5-service:
    build: ./yolov5-app  # 指向含main.py和修改后Dockerfile的YoloV5项目目录
    networks:
      - django-yolo-network
    ports:
      - "8001:8000"
    # 若需GPU支持，添加以下配置（需安装NVIDIA Docker）
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    depends_on:
      - postgres

volumes:
  postgres_data:

关键注意事项 #

• GPU支持：若YoloV5需GPU加速，确保主机安装NVIDIA驱动和NVIDIA Docker，并在Compose中配置GPU资源。
• 环境变量：用环境变量管理YoloV5 API地址，方便不同环境切换。
• 错误处理：Django中需处理API调用异常（如YoloV5容器宕机、网络超时），返回友好错误信息。
• 性能优化：生产环境中，YoloV5 API可启用uvicorn多进程模式（如uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4），Django用gunicorn多进程运行。

内容的提问来源于stack exchange，提问作者Joker King