How to use yolov5 model in django

Question

I want to use the yolov5 model in django, but I got trouble.

What I want to make is if user upload an image to django server, object detection made by with the yolov5 model, and then the result displayed on the web. The process itself is simple. But I am not sure how to attach yolov5 api and django.

As those of you who have used yolo know, yolo basically runs apis based on commands.

!python train.py --img 512 --batch 16 --epochs 100 --data ~~   # for training
!python detect.py --weights'/content/yolov5/runs/~~            # for detection

Is there any way to run this in a Django view?

I tried python functions such as execfile() and os.system() that execute commands in an interpreter way within Python script, but it didn't work properly. (I think the paths between Django and Yolo were twisted.)

In fact, if possible, it is best to load yolo api as a module like Keras and run it like a function, rather than the command method. But I can't find a way to use yolov5 like a module or something like that.

How can I solve this problem? Please give me some advice.

Answer 1

Object detection using Django and Yolov5

Configure your Django project with yolov5:

I have uploaded the django-object-detection here, for easy reference.

Models.py:

import os

from django.db import models
from django.utils.translation import gettext_lazy as _


class ImageModel(models.Model):
    image = models.ImageField(_("image"), upload_to='images')

    class Meta:
        verbose_name = "Image"
        verbose_name_plural = "Images"

    def __str__(self):
        return str(os.path.split(self.image.path)[-1])

Views.py:

import io
from PIL import Image as im
import torch

from django.shortcuts import render
from django.views.generic.edit import CreateView

from .models import ImageModel
from .forms import ImageUploadForm


class UploadImage(CreateView):
    model = ImageModel
    template_name = 'image/imagemodel_form.html'
    fields = ["image"]

    def post(self, request, *args, **kwargs):
        form = ImageUploadForm(request.POST, request.FILES)
        if form.is_valid():
            img = request.FILES.get('image')
            img_instance = ImageModel(
                image=img
            )
            img_instance.save()

            uploaded_img_qs = ImageModel.objects.filter().last()
            img_bytes = uploaded_img_qs.image.read()
            img = im.open(io.BytesIO(img_bytes))

            # Change this to the correct path
            path_hubconfig = "absolute/path/to/yolov5_code"
            path_weightfile = "absolute/path/to/yolov5s.pt"  # or any custom trained model

            model = torch.hub.load(path_hubconfig, 'custom',
                               path=path_weightfile, source='local')

            results = model(img, size=640)
            results.render()
            for img in results.imgs:
                img_base64 = im.fromarray(img)
                img_base64.save("media/yolo_out/image0.jpg", format="JPEG")

            inference_img = "/media/yolo_out/image0.jpg"

            form = ImageUploadForm()
            context = {
                "form": form,
                "inference_img": inference_img
            }
            return render(request, 'image/imagemodel_form.html', context)

        else:
            form = ImageUploadForm()
        context = {
            "form": form
        }
        return render(request, 'image/imagemodel_form.html', context)

A lot can be modified here in this views.py, like adding more functions and logic, but the aim here is to connect yolov5 & Django. Configurations inside views.py are most important as it is the gateway to the yolov5 hubconf.py file.

Forms.py

from django import forms
from .models import ImageModel


class ImageUploadForm(forms.ModelForm):
    class Meta:
        model = ImageModel
        fields = ['image']

Imagemodel_form.html

{% extends "base.html" %}
{% load crispy_forms_tags %}


{% block leftbar %}
    <div class="col-sm-3">

    </div>
{% endblock leftbar %} 


{% block content %}
    <div class="col-sm-9">
        <div id="uploadedImage"></div>
            <div class="mt-4">
                <form action="" enctype="multipart/form-data" id="imageUploadForm" method="post"> 
                    {% csrf_token %}
                    {{ form|crispy }}
                    <button class="btn btn-outline-success" type="submit">Submit</button>
                </form>
            </div>
        </div>
        <div class="mt-4">
            {% if inference_img %}
                <img src="{{inference_img}}" class="img-fluid" />
            {% else %}
                The inferenced image will be displayed here.
            {% endif %}
        </div>    
    </div>
{% endblock content %}

Simple webpage initially:

After detection: