Plankton_Detector/PlanktonDetector/DetectionApp/views.py

from typing import Any
from django.db.models.query import QuerySet
from django.http import HttpResponse
from django.shortcuts import render
from django.views import View
from .forms import DetectForm
from django.views.generic import ListView, DetailView
from django.contrib.auth.decorators import login_required
from django.utils.decorators import method_decorator
from .models import PredictionBatch, UploadImage, PredictedImage
from .utils import predict_image
from django.contrib.auth.mixins import LoginRequiredMixin


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})

    @method_decorator(login_required)
    def post(self, request, *args, **kwargs):
        form = self.form_class(request.POST, request.FILES)
        files = request.FILES.getlist("image")
        predictions = []
        if form.is_valid():
            pred_batch = PredictionBatch.objects.create(owner=request.user)
            for f in files:
                image = UploadImage(
                    image=f,
                )
                image.save()
                prediciton_results = predict_image(image)
                image.predicted_image_url = f"{image.image.name.split('.')[0]}_predicted.{image.image.name.split('.')[-1]}"
                image.save()
                try:
                    results_metrics = prediciton_results
                except IndexError as e:
                    predicted_image = PredictedImage.objects.create(
                        original_image=image,
                        image=image.predicted_image_url,
                        prediction_data={"data": "no predicitions"},
                    )
                else:
                    predicted_image = PredictedImage.objects.create(
                        original_image=image,
                        image=image.predicted_image_url,
                        prediction_data=results_metrics,
                    )
                predictions.append(predicted_image)
            pred_batch.images.add(*predictions)
            pred_batch.save()
            return render(
                request,
                "results.html",
                {
                    "img_saved": True,
                    "img": pred_batch,
                },
            )
        else:
            return render(request, "upload.html", {"form": form})


class ListHistory(LoginRequiredMixin, ListView):
    model = PredictionBatch
    queryset = PredictionBatch.objects.all()
    template_name = "history.html"
    paginate_by = 3

    def get_queryset(self) -> QuerySet[Any]:
        queryset = PredictionBatch.objects.filter(owner=self.request.user).order_by(
            "-date_predicted"
        )
        return queryset


class DetectionDetails(LoginRequiredMixin, DetailView):
    model = PredictionBatch
    template_name = "results.html"
    context_object_name = "img"


def download_pred_res(request, pk):
    pred_batch = PredictionBatch.objects.get(pk=pk)
    response = HttpResponse(content_type="text/plain")
    response["Content-Disposition"] = "attachment; filename=predictions.txt"
    for img in pred_batch.images.all():
        response.write(img.prediction_data)
    return response