update: remove image resize

app.py

@@ -1,171 +1,170 @@
-import torch
-from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
-import os
-import gradio as gr
-import numpy as np
-from PIL import Image
-from PIL.ImageOps import grayscale
-import cv2
-import torch
-import gc
-import math
-import cvzone
-from cvzone.PoseModule import PoseDetector
-from cvzone.FaceMeshModule import FaceMeshDetector
-import spaces
-
-choker_images = [Image.open(os.path.join("short_necklace", x)) for x in os.listdir("short_necklace")]
-person_images = [Image.open(os.path.join("without_necklace", x)) for x in os.listdir("without_necklace")]
-
-model_id = "stabilityai/stable-diffusion-2-inpainting"
-pipeline = StableDiffusionInpaintPipeline.from_pretrained(
-    model_id, torch_dtype=torch.float16
-)
-pipeline = pipeline.to("cuda")
-
-detector = PoseDetector()
-meshDetector = FaceMeshDetector(staticMode=True, maxFaces=1)
-
-
-def clear_func():
-    torch.cuda.empty_cache()
-    gc.collect()
-
-
-@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)
-
-    image = detector.findPose(image)
-    lmList, _ = detector.findPosition(image, bboxWithHands=False, draw=False)
-
-    img, faces = 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(
-            detector.findAngle(
-                p1=(avg_x2, avg_y2), p2=(avg_x1, avg_y1), p3=(avg_x2, avg_y1)
-            )[0]
-        )
-    else:
-        angle = math.ceil(
-            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)
-    # 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
-
-    gc.collect()
-
-    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)
-
-    arr = arr_orig.copy()
-    mask_y = np.where(arr == arr[arr != 0][0])[0][0]
-    arr[mask_y:, :] = 255
-
-    new = Image.fromarray(arr)
-
-    mask = new.copy()
-
-    orig_size = image.size
-
-    image = image.resize((512, 512))
-    mask = mask.resize((512, 512))
-
-    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="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)
+import torch
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
+import os
+import gradio as gr
+import numpy as np
+from PIL import Image
+from PIL.ImageOps import grayscale
+import cv2
+import torch
+import gc
+import math
+import cvzone
+from cvzone.PoseModule import PoseDetector
+from cvzone.FaceMeshModule import FaceMeshDetector
+import spaces
+
+choker_images = [Image.open(os.path.join("short_necklace", x)) for x in os.listdir("short_necklace")]
+person_images = [Image.open(os.path.join("without_necklace", x)) for x in os.listdir("without_necklace")]
+
+model_id = "stabilityai/stable-diffusion-2-inpainting"
+pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+    model_id, torch_dtype=torch.float16
+)
+pipeline = pipeline.to("cuda")
+
+detector = PoseDetector()
+meshDetector = FaceMeshDetector(staticMode=True, maxFaces=1)
+
+
+def clear_func():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+@spaces.GPU
+def clothing_try_on_n_necklace_try_on(image, jewellery):
+    image = np.array(image)
+    copy_image = image.copy()
+    jewellery = np.array(jewellery)
+
+    image = detector.findPose(image)
+    lmList, _ = detector.findPosition(image, bboxWithHands=False, draw=False)
+
+    img, faces = 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(
+            detector.findAngle(
+                p1=(avg_x2, avg_y2), p2=(avg_x1, avg_y1), p3=(avg_x2, avg_y1)
+            )[0]
+        )
+    else:
+        angle = math.ceil(
+            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)
+    # 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
+
+    gc.collect()
+
+    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)
+
+    arr = arr_orig.copy()
+    mask_y = np.where(arr == arr[arr != 0][0])[0][0]
+    arr[mask_y:, :] = 255
+
+    new = Image.fromarray(arr)
+
+    mask = new.copy()
+
+    orig_size = image.size
+
+    image = image.resize((512, 512))
+    mask = mask.resize((512, 512))
+
+    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="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)