content_quality_tool/bg_remover/views.py

# background_remover/views.py
import io
from PIL import Image
from django.http import HttpResponse
from rest_framework.views import APIView
from rest_framework.parsers import MultiPartParser
from .services import BiRefNetService
from django.shortcuts import render
import zipfile
import time

class BackgroundRemovalView(APIView):
    parser_classes = [MultiPartParser]

    def get(self, request):
        return render(request, "bg_remover_index.html")

    def post(self, request, *args, **kwargs):
        if 'image' not in request.FILES:
            return HttpResponse("No image provided", status=400)

        # 1. Load image from request
        input_file = request.FILES['image']
        input_image = Image.open(input_file)

        # 2. Process via Service
        service = BiRefNetService()
        output_image = service.remove_background(input_image)

        # 3. Prepare response
        buffer = io.BytesIO()
        output_image.save(buffer, format="PNG")
        
        return HttpResponse(buffer.getvalue(), content_type="image/png")
    
class BulkBackgroundRemovalAPI(APIView):
    parser_classes = [MultiPartParser]

    def post(self, request):
        files = request.FILES.getlist("images")

        if not files:
            return HttpResponse({"error": "No images provided"}, status=400)

        service = BiRefNetService()
        start_time = time.time()

        zip_buffer = io.BytesIO()

        # with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
            # for file in files:
            #     try:
            #         image = Image.open(file).convert("RGB")
            #         output = service.remove_background(image)

            #         img_buffer = io.BytesIO()
            #         output.save(img_buffer, format="PNG")

            #         name = file.name.rsplit(".", 1)[0] + "_whitebg.png"
            #         zip_file.writestr(name, img_buffer.getvalue())

            #     except Exception as e:
            #         print(f"Error processing {file.name}: {e}")

        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as master_zip:
            for file in files:
                # CHECK IF FILE IS A ZIP
                if file.name.lower().endswith('.zip'):
                    with zipfile.ZipFile(file, 'r') as user_zip:
                        for inner_file_name in user_zip.namelist():
                            # Skip directories and non-image files
                            if inner_file_name.endswith(('.png', '.jpg', '.jpeg', '.webp')):
                                with user_zip.open(inner_file_name) as inner_file:
                                    img_data = inner_file.read()
                                    image = Image.open(io.BytesIO(img_data)).convert("RGB")
                                    output = service.remove_background(image)
                                    
                                    img_io = io.BytesIO()
                                    output.save(img_io, format="PNG")
                                    
                                    # Create clean name for the output zip
                                    clean_name = inner_file_name.rsplit(".", 1)[0].split('/')[-1] + "_no_bg.png"
                                    master_zip.writestr(clean_name, img_io.getvalue())
                
                # PROCESS INDIVIDUAL IMAGES (Existing logic)
                else:
                    try:
                        image = Image.open(file).convert("RGB")
                        output = service.remove_background(image)
                        img_io = io.BytesIO()
                        output.save(img_io, format="PNG")
                        name = file.name.rsplit(".", 1)[0] + "_no_bg.png"
                        master_zip.writestr(name, img_io.getvalue())
                    except Exception as e:
                        print(f"Error processing {file.name}: {e}")
        zip_buffer.seek(0)
        total_time = round(time.time() - start_time, 2)

        response = HttpResponse(
            zip_buffer.getvalue(),
            content_type="application/zip"
        )
        response["Content-Disposition"] = "attachment; filename=processed_images.zip"
        response["X-Processing-Time"] = f"{total_time}s"

        return response