Plankton_Detector/PlanktonDetector/DetectionApp/views.py

from django.shortcuts import render, redirect
from django.views import View
from .forms import DetectForm
from django.conf import settings
from ultralytics import YOLO
from django.contrib.auth.forms import AuthenticationForm, UserCreationForm
from django.contrib.auth import authenticate, login, logout

# Create your views here.
MODEL = YOLO("best.pt")


class DetectView(View):
    form_class = DetectForm
    template_name = "upload.html"

    def get(self, request, *args, **kwargs):
        form = self.form_class()
        return render(request, "upload.html", {"form": form})

    def post(self, request, *args, **kwargs):
        form = self.form_class(request.POST, request.FILES)
        if form.is_valid():
            image = form.save()
            MODEL.predict(
                image.image.path,
                save=True,
                project=settings.MEDIA_ROOT,
                name=f"{image.image.name}_predicted",
                imgsz=[640, 640],
            )
            saved = form.is_valid()
            return render(
                request,
                "upload.html",
                {
                    "img_saved": saved,
                    "img_path": f"{image.image.name}_predicted/{image.image.name.split('/')[-1]}",
                },
            )
        else:
            return render(request, "upload.html", {"form": form})


class LoginView(View):
    form_class = AuthenticationForm
    template_name = "login_signup.html"

    def get(self, request, *args, **kwargs):
        form = self.form_class()
        return render(request, self.template_name, {"form": form})

    def post(self, request, *args, **kwargs):
        form = self.form_class(data=request.POST)
        if form.is_valid():
            user = authenticate(
                username=form.cleaned_data["username"],
                password=form.cleaned_data["password"],
            )
            if user is not None:
                login(request, user)
                return redirect("detect")
        return render(request, self.template_name, {"form": form})


class SignupView(View):
    form_class = UserCreationForm
    template_name = "login_signup.html"

    def get(self, request, *args, **kwargs):
        form = self.form_class()
        return render(request, self.template_name, {"form": form, "signup": True})

    def post(self, request, *args, **kwargs):
        form = self.form_class(data=request.POST)
        if form.is_valid():
            user = form.save()
            login(request, user)
            return redirect("detect")
        return render(request, self.template_name, {"form": form, "signup": True})


def logout_view(request):
    logout(request)
    return redirect("detect")