update: remove necklacetryon

app.py

@@ -1,17 +1,12 @@
-import torch
 import cv2
 import torch
 import gc
-import math
-import cvzone
 import os
 import spaces
 import gradio as gr
 import numpy as np
 from PIL import Image
 from PIL.ImageOps import grayscale
-from cvzone.PoseModule import PoseDetector
-from cvzone.FaceMeshModule import FaceMeshDetector
 from diffusers import StableDiffusionInpaintPipeline
 
 
@@ -26,86 +21,13 @@ class NecklaceTryOn:
         )
         self.pipeline = self.pipeline.to("cuda")
 
-        self.detector = PoseDetector()
-        self.meshDetector = FaceMeshDetector(staticMode=True, maxFaces=1)
-
     def clear_func(self):
         torch.cuda.empty_cache()
         gc.collect()
 
     @spaces.GPU
-    def apply_necklace(self, image, jewellery):
-        """Apply necklace on the image and return modified image and mask."""
-        image = np.array(image)
-        copy_image = image.copy()
-        jewellery = np.array(jewellery)
-
-        image = self.detector.findPose(image)
-        lmList, _ = self.detector.findPosition(image, bboxWithHands=False, draw=False)
-
-        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
-            else:
-                continue
-        offset = int(0.8 * xdist * (offset_orig / 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 = cvzone.overlayPNG(copy_image, jewellery, (avg_x1, y_coordinate))
-
-        blackedNecklace = np.zeros(shape=copy_image.shape)
-        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
-
-        gc.collect()
-
-        return Image.fromarray(result.astype(np.uint8)), Image.fromarray(binaryMask.astype(np.uint8)).convert("RGB")
-
-    @spaces.GPU
-    def clothing_try_on_n_necklace_try_on(self, image, jewellery):
+    def clothing_try_on_n_necklace_try_on(self, result_image, mask):
         """Main method for clothing and necklace try-on."""
-        result_image, mask = self.apply_necklace(image, jewellery)
-
         jewellery_mask = Image.fromarray(
             np.bitwise_and(np.array(mask), np.array(result_image))
         )
@@ -155,16 +77,12 @@ class NecklaceTryOn:
         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)
+                mask_image = gr.Image(label="Mask Image", type="pil", image_mode="RGB", interactive=True)
                 outputOne = gr.Image(label="Output", interactive=False)
 
-            with gr.Row():
-                gr.Examples(examples=self.choker_images, inputs=[selectedNecklace], label="Select Necklace")
-                gr.Examples(examples=self.person_images, inputs=[inputImage], label="Select Model")
-
             submit = gr.Button("Apply")
 
-            submit.click(fn=self.clothing_try_on_n_necklace_try_on, inputs=[inputImage, selectedNecklace],
+            submit.click(fn=self.clothing_try_on_n_necklace_try_on, inputs=[inputImage, mask_image],
                          outputs=[outputOne])
 
         interface.launch(debug=True)