mannequin_to_model.py

"""
Created By: ishwor subedi
Date: 2024-07-03
"""
import os
import cv2
import base64
import requests
import numpy as np
from io import BytesIO
from PIL import Image
from fastapi import File, UploadFile, Form, Depends
from fastapi import APIRouter
from fastapi.responses import JSONResponse
import json
from src.pipeline.main_pipeline import MainPipeline
from supabase import create_client

secure_router = APIRouter()

pipeline = MainPipeline()

SUPABASE_KEY = os.getenv("SUPABASE_KEY")
SUPABASE_URL = os.getenv("SUPABASE_URL")

supabase = create_client(SUPABASE_URL, SUPABASE_KEY)


def read_return(url):
    res = requests.get(url)
    return res.content


@secure_router.post("/mannequin_to_model")
async def mannequinToModel(
        store_name: str = Form(...),
        clothing_category: str = Form(...),
        product_id: str = Form(...),
        body_structure: str = Form(...),
        skin_complexion: str = Form(...),
        facial_structure: str = Form(...),
        person_img: UploadFile = File(...)
):
    if body_structure == "medium":
        body_structure = "fat"

    person_image = await person_img.read()

    mannequin_image_url = f"{SUPABASE_URL}/storage/v1/object/public/ClothingTryOn/{store_name}/{clothing_category}/{product_id}/{product_id}_{skin_complexion}_{facial_structure}_{body_structure}.webp"
    mannequin_img = read_return(mannequin_image_url)

    try:
        person_image, mannequin_image = Image.open(BytesIO(person_image)), Image.open(BytesIO(mannequin_img))
    except:

        error_message = {
            "code": 404,
            "error": "The requested resource is not available. Please verify the availability and try again."

        }

        return JSONResponse(content=error_message, status_code=404)

    mannequin_image = cv2.cvtColor(np.array(mannequin_image), cv2.COLOR_RGB2BGR)
    person_image = cv2.cvtColor(np.array(person_image), cv2.COLOR_RGB2BGR)
    result = pipeline.face_swap(mannequin_image, person_image)
    result = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
    inMemFile = BytesIO()
    result.save(inMemFile, format="WEBP", quality=85)
    outputBytes = inMemFile.getvalue()
    response = {
        "code": 200,
        "output": f"data:image/WEBP;base64,{base64.b64encode(outputBytes).decode('utf-8')}",
    }

    return response


@secure_router.post("/swap_image")
async def swapImage(image_one: UploadFile = File(...), ref_face: UploadFile = File(...)):
    image_one, ref_face = await image_one.read(), await ref_face.read()
    image_one_, ref_face_ = Image.open(BytesIO(image_one)), Image.open(BytesIO(ref_face))

    mannequin_image = cv2.cvtColor(np.array(image_one_), cv2.COLOR_RGB2BGR)
    person_image = cv2.cvtColor(np.array(ref_face_), cv2.COLOR_RGB2BGR)
    result = pipeline.face_swap(mannequin_image, person_image)
    result = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
    inMemFile = BytesIO()
    result.save(inMemFile, format="WEBP", quality=85)
    outputBytes = inMemFile.getvalue()
    response = {
        "code": 200,
        "output": f"data:image/WEBP;base64,{base64.b64encode(outputBytes).decode('utf-8')}",
    }
    return response


@secure_router.get("/mannequin_catalogue")
async def returnJsonData(gender: str):
    folderImageURL = supabase.storage.get_bucket("JSON").create_signed_url(
        path=os.path.join("MannequinInfo.json"), expires_in=3600)["signedURL"]
    r = requests.get(folderImageURL).content.decode()

    mannequin_data = json.loads(r)

    if gender.lower() == "female":
        res = [item for item in mannequin_data if item["gender"] == "female"]
    elif gender.lower() == "male":
        res = [item for item in mannequin_data if item["gender"] == "male"]
    else:
        res = []

    return res