update: combine inpaint and necklacetryon

.gitignore

@@ -0,0 +1 @@
+.idea

app.py

@@ -14,10 +14,7 @@ import cvzone
 from cvzone.PoseModule import PoseDetector
 import spaces
 
-# images
 choker_images = [Image.open(os.path.join("choker", x)) for x in os.listdir("choker")]
-short_necklaces = [Image.open(os.path.join("short_necklace", x)) for x in os.listdir("short_necklace")]
-long_necklaces = [Image.open(os.path.join("long_haram", x)) for x in os.listdir("long_haram")]
 person_images = [Image.open(os.path.join("without_necklace", x)) for x in os.listdir("without_necklace")]
 
 # initialising the stable diffusion model
@@ -27,14 +24,16 @@ pipeline = StableDiffusionInpaintPipeline.from_pretrained(
 )
 pipeline = pipeline.to("cuda")
 
+
 # functions
 def clearFunc():
     torch.cuda.empty_cache()
     gc.collect()
 
-def necklaceTryOnPipeline(image, jewellery):
-    global binaryMask
-    
+
+# SD Model
+@spaces.GPU
+def sd_inpaint(image, jewellery):
     image = np.array(image)
     copy_image = image.copy()
     jewellery = np.array(jewellery)
@@ -93,27 +92,22 @@ def necklaceTryOnPipeline(image, jewellery):
     available_space = copy_image.shape[0] - y_coordinate
     extra = jewellery.shape[0] - available_space
     if extra > 0:
-        jewellery = jewellery[extra + 10 :, :]
+        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)
+        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
-        return Image.fromarray(result.astype(np.uint8)), Image.fromarray(binaryMask.astype(np.uint8))
-
-# SD Model
-@spaces.GPU
-def sd_inpaint(image, mask):
-#    image = Image.fromarray(image)
-#    mask = Image.fromarray(mask)
+        image = Image.fromarray(result.astype(np.uint8))
+        mask = Image.fromarray(binaryMask.astype(np.uint8))
 
     jewellery_mask = Image.fromarray(
         np.bitwise_and(np.array(mask), np.array(image))
@@ -137,57 +131,47 @@ def sd_inpaint(image, mask):
     mask = mask.resize((512, 512))
 
     results = []
-    for colour in ["Red", "Blue", "Green"]:
-        prompt = f"{colour}, 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)
+    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
     ]
     clearFunc()
-    return results[0], results[1], results[2]
+    return results[0]
+
 
 # interface
 
 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)
-    hiddenMask = gr.Image(visible = False, type = "pil")
-  with gr.Row():
-    gr.Examples(examples = person_images, inputs=[inputImage], label="Models")
-  with gr.Row():
-    gr.Examples(examples = choker_images, inputs = [selectedNecklace], label = "Chokers")
-  with gr.Row():
-    gr.Examples(examples = short_necklaces, inputs = [selectedNecklace], label = "Short Necklaces")
-  with gr.Row():
-    gr.Examples(examples = long_necklaces, inputs = [selectedNecklace], label = "Long Necklaces")
-  with gr.Row():
-    outputOne = gr.Image(label = "Output 1", interactive = False)
-    outputTwo = gr.Image(label = "Output 2", interactive = False)
-    outputThree = gr.Image(label = "Output 3", interactive = False)
-  with gr.Row():
-    submit = gr.Button("Enter")
-
-  selectedNecklace.change(fn = necklaceTryOnPipeline, inputs = [inputImage, selectedNecklace], outputs = [inputImage, hiddenMask])
-  submit.click(fn = sd_inpaint, inputs = [inputImage, hiddenMask], outputs = [outputOne, outputTwo, outputThree])
-
-
-interface.launch(debug = True)
+    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)
+    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")
+
+    submit.click(fn=sd_inpaint, inputs=[inputImage, selectedNecklace], outputs=[outputOne, outputTwo, outputThree])
+
+interface.launch(debug=True)