update: necklace placement logic and gradio ui

app.py

@@ -22,6 +22,9 @@ pipeline = StableDiffusionInpaintPipeline.from_pretrained(
 )
 pipeline = pipeline.to("cuda")
 
+detector = PoseDetector()
+meshDetector = FaceMeshDetector(staticMode=True, maxFaces=1)
+
 
 def clear_func():
     torch.cuda.empty_cache()
@@ -30,29 +33,33 @@ def clear_func():
 
 @spaces.GPU
 def clothing_try_on_n_necklace_try_on(image, jewellery):
+    image, jewellery = image.convert("RGB").resize((3000, 3000)), jewellery.convert("RGBA")
     image = np.array(image)
     copy_image = image.copy()
     jewellery = np.array(jewellery)
 
-    detector = PoseDetector()
+    logger.info(f"NECKLACE TRY ON :: detecting pose and landmarks :: {storename}")
 
     image = detector.findPose(image)
-    lmList, bBoxInfo = detector.findPosition(image, bboxWithHands=False, draw=False)
+    lmList, _ = detector.findPosition(image, bboxWithHands=False, draw=False)
 
-    pt12, pt11, pt10, pt9 = (
-        lmList[12][:2],
-        lmList[11][:2],
-        lmList[10][:2],
-        lmList[9][:2],
-    )
+    img, faces = meshDetector.findFaceMesh(image, draw=False)
+    leftLandmarkIndex = 172
+    rightLandmarkIndex = 397
 
-    avg_x1 = int(pt12[0] + (pt10[0] - pt12[0]) / 2)
-    avg_y1 = int(pt12[1] - (pt12[1] - pt10[1]) / 2)
+    leftLandmark, rightLandmark = faces[0][leftLandmarkIndex], faces[0][rightLandmarkIndex]
+    landmarksDistance = int(
+        ((leftLandmark[0] - rightLandmark[0]) ** 2 + (leftLandmark[1] - rightLandmark[1]) ** 2) ** 0.5)
 
-    avg_x2 = int(pt11[0] - (pt11[0] - pt9[0]) / 2)
-    avg_y2 = int(pt11[1] - (pt11[1] - pt9[1]) / 2)
+    logger.info(f"NECKLACE TRY ON :: estimating neck points")
 
-    image_gray = cv2.cvtColor(jewellery, cv2.COLOR_BGRA2GRAY)
+    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 ")
 
     if avg_y2 < avg_y1:
         angle = math.ceil(
@@ -72,100 +79,101 @@ def clothing_try_on_n_necklace_try_on(image, jewellery):
     origImgRatio = xdist / jewellery.shape[1]
     ydist = jewellery.shape[0] * origImgRatio
 
+    logger.info(f"NECKLACE TRY ON :: adding offset based on the necklace shape")
+
+    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
         else:
             continue
-
     offset = int(0.8 * xdist * (offset_orig / jewellery.shape[1]))
     jewellery = cv2.resize(
-        jewellery, (int(xdist), int(ydist)), interpolation=cv2.INTER_CUBIC
+        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
-    if extra > 0:
-        jewellery = jewellery[extra + 10:, :]
-        return necklaceTryOnPipeline(
-            Image.fromarray(copy_image), Image.fromarray(jewellery)
-        )
-    else:
-        result = cvzone.overlayPNG(copy_image, jewellery, (avg_x1, y_coordinate))
-        # masking
-        blackedNecklace = np.zeros(shape=copy_image.shape)
-        # overlay
-        cvzone.overlayPNG(blackedNecklace, jewellery, (avg_x1, y_coordinate))
-        blackedNecklace = cv2.cvtColor(blackedNecklace.astype(np.uint8), cv2.COLOR_BGR2GRAY)
-        binaryMask = blackedNecklace * ((blackedNecklace > 5) * 255)
-        binaryMask[binaryMask >= 255] = 255
-        binaryMask[binaryMask < 255] = 0
-        image = Image.fromarray(result.astype(np.uint8))
-        mask = Image.fromarray(binaryMask.astype(np.uint8)).convert("RGB")
-
-        jewellery_mask = Image.fromarray(
-            np.bitwise_and(np.array(mask), np.array(image))
-        )
-        arr_orig = np.array(grayscale(mask))
 
-        image = cv2.inpaint(np.array(image), arr_orig, 15, cv2.INPAINT_TELEA)
-        image = Image.fromarray(image)
+    logger.info(f"NECKLACE TRY ON :: generating output :: {storename}")
+
+    result = cvzone.overlayPNG(copy_image, jewellery, (avg_x1, y_coordinate))
+
+    blackedNecklace = np.zeros(shape=copy_image.shape)
+    # overlay
+    cvzone.overlayPNG(blackedNecklace, jewellery, (avg_x1, y_coordinate))
+    blackedNecklace = cv2.cvtColor(blackedNecklace.astype(np.uint8), cv2.COLOR_BGR2GRAY)
+    binaryMask = blackedNecklace * ((blackedNecklace > 5) * 255)
+    binaryMask[binaryMask >= 255] = 255
+    binaryMask[binaryMask < 255] = 0
 
-        arr = arr_orig.copy()
-        mask_y = np.where(arr == arr[arr != 0][0])[0][0]
-        arr[mask_y:, :] = 255
+    gc.collect()
+
+    image = Image.fromarray(result.astype(np.uint8))
+    mask = Image.fromarray(binaryMask.astype(np.uint8)).convert("RGB")
 
-        new = Image.fromarray(arr)
+    jewellery_mask = Image.fromarray(
+        np.bitwise_and(np.array(mask), np.array(image))
+    )
+    arr_orig = np.array(grayscale(mask))
 
-        mask = new.copy()
+    image = cv2.inpaint(np.array(image), arr_orig, 15, cv2.INPAINT_TELEA)
+    image = Image.fromarray(image)
 
-        orig_size = image.size
+    arr = arr_orig.copy()
+    mask_y = np.where(arr == arr[arr != 0][0])[0][0]
+    arr[mask_y:, :] = 255
 
-        image = image.resize((512, 512))
-        mask = mask.resize((512, 512))
+    new = Image.fromarray(arr)
 
-        results = []
-        prompt = f" South Indian Saree, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple"
-        negative_prompt = "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly"
+    mask = new.copy()
 
-        output = pipeline(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=image,
-            mask_image=mask,
-            strength=0.95,
-            guidance_score=9,
-            # generator = torch.Generator("cuda").manual_seed(42)
-        ).images[0]
+    orig_size = image.size
 
-        output = output.resize(orig_size)
-        temp_generated = np.bitwise_and(
-            np.array(output),
-            np.bitwise_not(np.array(Image.fromarray(arr_orig).convert("RGB"))),
-        )
-        results.append(temp_generated)
+    image = image.resize((512, 512))
+    mask = mask.resize((512, 512))
 
-        results = [
-            Image.fromarray(np.bitwise_or(x, np.array(jewellery_mask))) for x in results
-        ]
-        clear_func()
-        return results[0]
+    results = []
+    prompt = f" South Indian Saree, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple"
+    negative_prompt = "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly"
+
+    output = pipeline(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image,
+        mask_image=mask,
+        strength=0.95,
+        guidance_score=9,
+        # generator = torch.Generator("cuda").manual_seed(42)
+    ).images[0]
+
+    output = output.resize(orig_size)
+    temp_generated = np.bitwise_and(
+        np.array(output),
+        np.bitwise_not(np.array(Image.fromarray(arr_orig).convert("RGB"))),
+    )
+    results.append(temp_generated)
+
+    results = [
+        Image.fromarray(np.bitwise_or(x, np.array(jewellery_mask))) for x in results
+    ]
+    clear_func()
+    return results[0]
 
 
 with gr.Blocks() as interface:
     with gr.Row():
         inputImage = gr.Image(label="Input Image", type="pil", image_mode="RGB", interactive=True)
         selectedNecklace = gr.Image(label="Selected Necklace", type="pil", image_mode="RGBA", visible=False)
+        outputOne = gr.Image(label="Output", interactive=False)
+
     with gr.Row():
-        gr.Examples(examples=choker_images, inputs=[selectedNecklace], label="Necklaces")
-        gr.Examples(examples=person_images, inputs=[inputImage], label="Models")
-    with gr.Row():
-        outputOne = gr.Image(label="Output 1", interactive=False)
-    with gr.Row():
-        submit = gr.Button("Enter")
+        gr.Examples(examples=choker_images, inputs=[selectedNecklace], label="Select Necklace")
+        gr.Examples(examples=person_images, inputs=[inputImage], label="Select Model")
+
+    submit = gr.Button("Apply")
 
     submit.click(fn=clothing_try_on_n_necklace_try_on, inputs=[inputImage, selectedNecklace], outputs=[outputOne])
 
 interface.launch(debug=True)
+