Spaces:
Runtime error
Runtime error
| from src.utils.exceptions import CustomException | |
| from cvzone.PoseModule import PoseDetector | |
| from src.utils.functions import getConfig | |
| from src.utils.logger import logger | |
| from PIL import Image | |
| import numpy as np | |
| import cvzone | |
| import math | |
| import cv2 | |
| class NecklaceTryOn: | |
| """ | |
| A class for simulating the wearing of necklaces in images. | |
| This class utilizes a pose detection algorithm to accurately overlay | |
| a necklace image onto a user's photo, adjusting for the user's neck | |
| position and orientation. | |
| Attributes: | |
| detector (PoseDetector): An instance of the PoseDetector for identifying | |
| body landmarks in images. | |
| config (ConfigParser): Configuration settings loaded from a specified | |
| configuration file (config.ini). | |
| Methods: | |
| necklaceTryOn(image: Image.Image, jewellery: Image.Image) -> Image.Image: | |
| Overlays a necklace onto the user's image based on detected pose | |
| landmarks and returns the resulting image. | |
| """ | |
| def __init__(self): | |
| """Initialize the NecklaceTryOn class with a PoseDetector and configuration settings.""" | |
| self.detector = PoseDetector() | |
| self.config = getConfig("config.ini") | |
| def necklaceTryOn(self, image: Image.Image, jewellery: Image.Image) -> Image.Image: | |
| """ | |
| Overlay a jewelry image onto a person's image to simulate wearing the jewelry. | |
| Args: | |
| image (Image.Image): The user's image, ideally captured in a standing, upright position. | |
| jewellery (Image.Image): The image of the jewelry piece (e.g., necklace) to be overlaid. | |
| Returns: | |
| Image.Image: A PIL Image depicting the user wearing the specified jewelry. | |
| Raises: | |
| CustomException: If an error occurs during the image processing. | |
| """ | |
| try: | |
| logger.info("converting images to numpy arrays") | |
| image = np.array(image) | |
| jewellery = np.array(jewellery) | |
| logger.info("creating a copy of original image for actual overlay") | |
| copyImage = image.copy() | |
| logger.info("detecting body landmarks from the input image") | |
| image = self.detector.findPose(image) | |
| lmList, _ = self.detector.findPosition(image, bboxWithHands = False, draw = False) | |
| pt12, pt11, pt10, pt9 = ( | |
| lmList[12][:2], | |
| lmList[11][:2], | |
| lmList[10][:2], | |
| lmList[9][:2], | |
| ) | |
| logger.info("calculating the precise neck points") | |
| avgX1 = int(pt12[0] + (pt10[0] - pt12[0]) / 1.75) | |
| avgY1 = int(pt12[1] - (pt12[1] - pt10[1]) / 1.75) | |
| avgX2 = int(pt11[0] - (pt11[0] - pt9[0]) / 1.75) | |
| avgY2 = int(pt11[1] - (pt11[1] - pt9[1]) / 1.75) | |
| logger.info("rescaling the necklace to appropriate dimensions") | |
| xDist = avgX2 - avgX1 | |
| origImgRatio = xDist / jewellery.shape[1] | |
| yDist = jewellery.shape[0] * origImgRatio | |
| jewellery = cv2.resize( | |
| jewellery, (int(xDist), int(yDist)), interpolation = cv2.INTER_CUBIC | |
| ) | |
| logger.info("calculating required offset to be added to the necklace image for perfect fitting") | |
| imageGray = cv2.cvtColor(jewellery, cv2.COLOR_BGRA2GRAY) | |
| for offsetOrig in range(imageGray.shape[1]): | |
| pixelValue = imageGray[0, :][offsetOrig] | |
| if (pixelValue != 255) & (pixelValue != 0): | |
| break | |
| else: | |
| continue | |
| offset = int(self.config.getfloat("NECKLACE TRY ON", "offsetFactor") * xDist * (offsetOrig / jewellery.shape[1])) | |
| yCoordinate = avgY1 - offset | |
| logger.info("tilting the necklace image as per the necklace points") | |
| angle = math.ceil( | |
| self.detector.findAngle( | |
| p1 = (avgX2, avgY2), p2 = (avgX1, avgY1), p3 = (avgX2, avgY1) | |
| )[0] | |
| ) | |
| if avgY2 < avgY1: | |
| pass | |
| else: | |
| angle = angle * -1 | |
| jewellery = cvzone.rotateImage(jewellery, angle) | |
| logger.info("checking if the necklace is getting out of the frame and trimming from above if needed") | |
| availableSpace = copyImage.shape[0] - yCoordinate | |
| extra = jewellery.shape[0] - availableSpace | |
| logger.info("applying the calculated settings") | |
| if extra > 0: | |
| jewellery = jewellery[extra + 10 :, :] | |
| return self.necklaceTryOn( | |
| Image.fromarray(copyImage), Image.fromarray(jewellery) | |
| ) | |
| else: | |
| result = cvzone.overlayPNG(copyImage, jewellery, (avgX1, yCoordinate)) | |
| result = Image.fromarray(result.astype(np.uint8)) | |
| return result | |
| except Exception as e: | |
| logger.error(CustomException(e)) | |
| print(CustomException(e)) |