86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
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")
|