src/components/necklaceTryOn.py

import time

from cvzone.FaceMeshModule import FaceMeshDetector
from src.utils import addWatermark, returnBytesData
from src.utils.exceptions import CustomException
from cvzone.PoseModule import PoseDetector
from src.utils.logger import logger
from dataclasses import dataclass
from typing import Union
from PIL import Image
import numpy as np
import cvzone
import math
import cv2
import gc


@dataclass
class NecklaceTryOnConfig:
    logoURL: str = "https://lvuhhlrkcuexzqtsbqyu.supabase.co/storage/v1/object/public/MagicMirror/FullImages/{}.png"


class NecklaceTryOn:
    def __init__(self) -> None:
        self.detector = PoseDetector()
        self.necklaceTryOnConfig = NecklaceTryOnConfig()
        self.meshDetector = FaceMeshDetector(staticMode=True, maxFaces=1)
        self.logo_cache = {}

    def necklaceTryOn(self, image: Image.Image, jewellery: Image.Image, storename: str) -> list[
        Union[Image.Image, str]]:
        try:
            logger.info(f">>> NECKLACE TRY ON STARTED :: {storename} <<<")

            image = np.array(image.convert("RGB").resize((3000, 3000)))
            copy_image = image.copy()
            jewellery = np.array(jewellery.convert("RGBA"))

            logger.info(f"NECKLACE TRY ON :: detecting pose and landmarks :: {storename}")
            image = self.detector.findPose(image)
            lmList, _ = self.detector.findPosition(image, bboxWithHands=False, draw=False)

            img, faces = self.meshDetector.findFaceMesh(image, draw=False)
            leftLandmark, rightLandmark = faces[0][172], faces[0][397]

            landmarksDistance = np.linalg.norm(np.array(leftLandmark) - np.array(rightLandmark))

            logger.info(f"NECKLACE TRY ON :: estimating neck points :: {storename}")
            avg_x1 = int(leftLandmark[0] - landmarksDistance * 0.12)
            avg_x2 = int(rightLandmark[0] + landmarksDistance * 0.12)
            avg_y1 = int(leftLandmark[1] + landmarksDistance * 0.5)
            avg_y2 = int(rightLandmark[1] + landmarksDistance * 0.5)

            logger.info(f"NECKLACE TRY ON :: scaling the necklace image :: {storename}")
            angle = math.ceil(self.detector.findAngle((avg_x2, avg_y2), (avg_x1, avg_y1), (avg_x2, avg_y1))[0])
            if avg_y2 >= avg_y1:
                angle *= -1

            xdist = avg_x2 - avg_x1
            origImgRatio = xdist / jewellery.shape[1]
            ydist = jewellery.shape[0] * origImgRatio

            logger.info(f"NECKLACE TRY ON :: adding offset based on the necklace shape :: {storename}")
            image_gray = cv2.cvtColor(jewellery, cv2.COLOR_BGRA2GRAY)
            offset = int(0.8 * xdist * (np.argmax(image_gray[0, :] != 255) / jewellery.shape[1]))

            jewellery = cv2.resize(jewellery, (int(xdist), int(ydist)), interpolation=cv2.INTER_AREA)
            jewellery = cvzone.rotateImage(jewellery, angle)
            y_coordinate = avg_y1 - offset
            available_space = copy_image.shape[0] - y_coordinate
            extra = jewellery.shape[0] - available_space

            headerText = "To see more of the necklace, please step back slightly." if extra > 0 else "success"

            logger.info(f"NECKLACE TRY ON :: generating output :: {storename}")
            result = cvzone.overlayPNG(copy_image, jewellery, (avg_x1, y_coordinate))
            image = Image.fromarray(result.astype(np.uint8))

            if storename not in self.logo_cache:
                self.logo_cache[storename] = Image.open(
                    returnBytesData(url=self.necklaceTryOnConfig.logoURL.format(storename)))
            result = addWatermark(background=image, logo=self.logo_cache[storename])

            # Create binary mask
            blackedNecklace = np.zeros_like(copy_image)
            cvzone.overlayPNG(blackedNecklace, jewellery, (avg_x1, y_coordinate))
            binaryMask = cv2.cvtColor(blackedNecklace.astype(np.uint8), cv2.COLOR_BGR2GRAY)
            binaryMask = (binaryMask > 5).astype(np.uint8) * 255
            mask = Image.fromarray(binaryMask).convert("RGB")

            gc.collect()
            return [result, headerText, mask]

        except Exception as e:
            logger.error(f">>> NECKLACE TRY ON ERROR: {str(e)} <<<")
            logger.error(f"{CustomException(e)}:: {storename}")
            return [None, "error", None]

    def necklaceTryOnMannequin(self, image, jewellery) -> list[
        Union[Image.Image, str]]:
        try:
            image_np = np.array(image.convert("RGB"))
            jewellery = np.array(jewellery.convert("RGBA"))

            height = image_np.shape[0]
            middle_point = height // 3
            upper_half = image_np[:middle_point, :]
            lower_half = image_np[middle_point:, :]

            upper_half = cv2.resize(upper_half, (upper_half.shape[1] * 3, upper_half.shape[0] * 3))
            copy_upper = upper_half.copy()

            # Apply pose detection on upper half
            upper_half = self.detector.findPose(upper_half)
            lmList, _ = self.detector.findPosition(upper_half, bboxWithHands=False, draw=True)

            img, faces = self.meshDetector.findFaceMesh(upper_half, draw=True)

            if not faces:
                return Exception("No face detected in the image")

            leftLandmark, rightLandmark = faces[0][172], faces[0][397]

            landmarksDistance = np.linalg.norm(np.array(leftLandmark) - np.array(rightLandmark))

            avg_x1 = int(leftLandmark[0] - landmarksDistance * 0.12)
            avg_x2 = int(rightLandmark[0] + landmarksDistance * 0.12)
            avg_y1 = int(leftLandmark[1] + landmarksDistance * 0.5)
            avg_y2 = int(rightLandmark[1] + landmarksDistance * 0.5)

            angle = math.ceil(self.detector.findAngle((avg_x2, avg_y2), (avg_x1, avg_y1), (avg_x2, avg_y1))[0])
            if avg_y2 >= avg_y1:
                angle *= -1

            xdist = avg_x2 - avg_x1
            origImgRatio = xdist / jewellery.shape[1]
            ydist = jewellery.shape[0] * origImgRatio

            image_gray = cv2.cvtColor(jewellery, cv2.COLOR_BGRA2GRAY)
            offset = int(0.8 * xdist * (np.argmax(image_gray[0, :] != 255) / jewellery.shape[1]))

            jewellery = cv2.resize(jewellery, (int(xdist), int(ydist)), interpolation=cv2.INTER_AREA)
            jewellery = cvzone.rotateImage(jewellery, angle)

            y_coordinate = avg_y1 - offset

            result_upper = cvzone.overlayPNG(copy_upper, jewellery, (avg_x1, y_coordinate))

            final_result = cv2.resize(result_upper, (image_np.shape[1], image_np.shape[0] - middle_point))

            final_result = np.vstack((final_result, lower_half))

            gc.collect()
            return Image.fromarray(final_result.astype(np.uint8)), Image.fromarray(result_upper.astype(np.uint8))

        except Exception as e:
            print(f"Error: {e}")
            return e

    def shoulderPointMaskGeneration(self, image: Image.Image) -> tuple[Image.Image, tuple[int, int], tuple[int, int]]:
        try:
            image = np.array(image)
            copy_image = image.copy()

            logger.info("SHOULDER POINT MASK GENERATION :: detecting pose and landmarks")
            image = self.detector.findPose(image)
            lmList, _ = self.detector.findPosition(image, bboxWithHands=False, draw=False)

            img, faces = self.meshDetector.findFaceMesh(image, draw=False)
            leftLandmark, rightLandmark = faces[0][172], faces[0][397]

            landmarksDistance = np.linalg.norm(np.array(leftLandmark) - np.array(rightLandmark))

            logger.info("SHOULDER POINT MASK GENERATION :: estimating neck points")
            # Using the same point calculation logic as necklaceTryOn
            avg_x1 = int(leftLandmark[0] - landmarksDistance * 0.12)
            avg_x2 = int(rightLandmark[0] + landmarksDistance * 0.12)
            avg_y1 = int(leftLandmark[1] + landmarksDistance * 0.5)
            avg_y2 = int(rightLandmark[1] + landmarksDistance * 0.5)

            logger.info("SHOULDER POINT MASK GENERATION :: generating shoulder point mask")
            mask = np.zeros_like(image[:, :, 0])
            mask[avg_y1:, :] = 255
            pts = np.array([[0, 0], [image.shape[1], 0], [avg_x2, avg_y2], [avg_x1, avg_y1]], np.int32)
            pts = pts.reshape((-1, 1, 2))
            cv2.fillPoly(mask, [pts], 0)

            black_n_white_mask = np.zeros_like(image[:, :, 0])
            black_n_white_mask[avg_y1:, :] = 255
            logger.info("SHOULDER POINT MASK GENERATION :: mask generated successfully")

            left_point = (avg_x1, avg_y1)
            right_point = (avg_x2, avg_y2)

            return Image.fromarray(black_n_white_mask.astype(np.uint8)), left_point, right_point

        except Exception as e:
            logger.error(f"SHOULDER POINT MASK GENERATION ERROR: {str(e)}")
            raise CustomException(e)

    def canvasPoints(self, image: Image.Image, jewellery: Image.Image, storename: str) -> dict:
        try:
            logger.info(f">>> NECKLACE TRY ON STARTED :: {storename} <<<")

            # Reading the images
            image, jewellery = image.convert("RGB").resize((3000, 3000)), jewellery.convert("RGBA")
            image = np.array(image)
            copy_image = image.copy()
            jewellery = np.array(jewellery)

            # Pose detection
            logger.info(f"NECKLACE TRY ON :: detecting pose and landmarks :: {storename}")
            image = self.detector.findPose(image)
            lmList, _ = self.detector.findPosition(image, bboxWithHands=False, draw=False)

            # Face mesh detection
            img, faces = self.meshDetector.findFaceMesh(image, draw=False)
            leftLandmarkIndex = 172
            rightLandmarkIndex = 397
            leftLandmark, rightLandmark = faces[0][leftLandmarkIndex], faces[0][rightLandmarkIndex]
            landmarksDistance = int(
                ((leftLandmark[0] - rightLandmark[0]) ** 2 + (leftLandmark[1] - rightLandmark[1]) ** 2) ** 0.5)

            avg_x1 = int(leftLandmark[0] - landmarksDistance * 0.12)
            avg_x2 = int(rightLandmark[0] + landmarksDistance * 0.12)
            avg_y1 = int(leftLandmark[1] + landmarksDistance * 0.5)
            avg_y2 = int(rightLandmark[1] + landmarksDistance * 0.5)

            if avg_y2 < avg_y1:
                angle = math.ceil(
                    self.detector.findAngle(
                        p1=(avg_x2, avg_y2), p2=(avg_x1, avg_y1), p3=(avg_x2, avg_y1)
                    )[0]
                )
            else:
                angle = math.ceil(
                    self.detector.findAngle(
                        p1=(avg_x2, avg_y2), p2=(avg_x1, avg_y1), p3=(avg_x2, avg_y1)
                    )[0]
                )
                angle = angle * -1

            xdist = avg_x2 - avg_x1
            origImgRatio = xdist / jewellery.shape[1]
            ydist = jewellery.shape[0] * origImgRatio
            image_gray = cv2.cvtColor(jewellery, cv2.COLOR_BGRA2GRAY)

            for offset_orig in range(image_gray.shape[1]):
                pixel_value = image_gray[0, :][offset_orig]
                if (pixel_value != 255) & (pixel_value != 0):
                    break

            offset = int(0.8 * xdist * (offset_orig / jewellery.shape[1]))
            y_coordinate = avg_y1 - offset

            points = {
                "left_shoulder": {"x": leftLandmark[0], "y": leftLandmark[1]},
                "right_shoulder": {"x": rightLandmark[0], "y": rightLandmark[1]},
                "necklace_top_left": {"x": avg_x1, "y": avg_y1},
                "necklace_top_right": {"x": avg_x2, "y": avg_y2},
                "necklace_bottom_left": {"x": avg_x1, "y": y_coordinate},
                "necklace_bottom_right": {"x": avg_x2, "y": y_coordinate},
                "angle": angle
            }

            gc.collect()

            return {
                "points": points,
                "status": "success"
            }

        except Exception as e:
            logger.error(f"{CustomException(e)}:: {storename}")
            raise CustomException(e)

    def necklaceTryOnWithPoints(self, image: Image.Image, jewellery: Image.Image, storename: str,
                                left_point: tuple[int, int], right_point: tuple[int, int]) -> list[
        Union[Image.Image, str]]:
        try:
            logger.info(f">>> NECKLACE TRY ON WITH POINTS STARTED :: {storename} <<<")

            image = np.array(image.convert("RGB"))
            # .resize((3000, 3000)))
            copy_image = image.copy()
            jewellery = np.array(jewellery.convert("RGBA"))

            logger.info(f"NECKLACE TRY ON :: scaling the necklace image based on given points :: {storename}")

            avg_x1 = left_point[0]
            avg_x2 = right_point[0]
            avg_y1 = left_point[1]
            avg_y2 = right_point[1]

            angle = math.ceil(self.detector.findAngle((avg_x2, avg_y2), (avg_x1, avg_y1), (avg_x2, avg_y1))[0])
            if avg_y2 >= avg_y1:
                angle *= -1

            xdist = avg_x2 - avg_x1
            origImgRatio = xdist / jewellery.shape[1]
            ydist = jewellery.shape[0] * origImgRatio

            logger.info(f"NECKLACE TRY ON :: adding offset based on the necklace shape :: {storename}")
            image_gray = cv2.cvtColor(jewellery, cv2.COLOR_BGRA2GRAY)
            offset = int(0.8 * xdist * (np.argmax(image_gray[0, :] != 255) / jewellery.shape[1]))

            jewellery = cv2.resize(jewellery, (int(xdist), int(ydist)), interpolation=cv2.INTER_AREA)
            jewellery = cvzone.rotateImage(jewellery, angle)
            y_coordinate = avg_y1 - offset
            available_space = copy_image.shape[0] - y_coordinate
            extra = jewellery.shape[0] - available_space

            headerText = "To see more of the necklace, please step back slightly." if extra > 0 else "success"

            logger.info(f"NECKLACE TRY ON :: generating output with given points :: {storename}")
            result = cvzone.overlayPNG(copy_image, jewellery, (avg_x1, y_coordinate))
            image = Image.fromarray(result.astype(np.uint8))

            if storename not in self.logo_cache:
                self.logo_cache[storename] = Image.open(
                    returnBytesData(url=self.necklaceTryOnConfig.logoURL.format(storename)))
            result = addWatermark(background=image, logo=self.logo_cache[storename])

            # Create binary mask
            blackedNecklace = np.zeros_like(copy_image)
            cvzone.overlayPNG(blackedNecklace, jewellery, (avg_x1, y_coordinate))
            binaryMask = cv2.cvtColor(blackedNecklace.astype(np.uint8), cv2.COLOR_BGR2GRAY)
            binaryMask = (binaryMask > 5).astype(np.uint8) * 255
            mask = Image.fromarray(binaryMask).convert("RGB")

            gc.collect()
            return [result, headerText, mask]

        except Exception as e:
            logger.error(f"{CustomException(e)}:: {storename}")
            raise CustomException(e)