commit

Former-commit-id: 0e58c57fa85967602d1f836634abad55d585c997

logs/app_debug.log

@@ -1371,3 +1371,21 @@ speed: {'preprocess': 1.9655227661132812, 'inference': 86.20810508728027, 'postp
 2024-03-22 06:39:26,709 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 06:43:26,086 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 06:53:54,562 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:03:17,990 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:09:38,955 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:17:59,975 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:28:54,339 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:41:45,300 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:50:46,880 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:05:52,674 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:11:07,093 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:49:32,092 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:56:43,084 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:13:03,681 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:23:45,335 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:29:54,960 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:31:28,680 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:39:32,351 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:47:05,180 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:51:28,523 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:53:18,039 [INFO] models - Kandinsky Inpainting Inference

logs/app_info.log

@@ -1371,3 +1371,21 @@ speed: {'preprocess': 1.9655227661132812, 'inference': 86.20810508728027, 'postp
 2024-03-22 06:39:26,709 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 06:43:26,086 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 06:53:54,562 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:03:17,990 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:09:38,955 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:17:59,975 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:28:54,339 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:41:45,300 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 07:50:46,880 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:05:52,674 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:11:07,093 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:49:32,092 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 08:56:43,084 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:13:03,681 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:23:45,335 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:29:54,960 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:31:28,680 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:39:32,351 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:47:05,180 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:51:28,523 [INFO] models - Kandinsky Inpainting Inference
+2024-03-22 09:53:18,039 [INFO] models - Kandinsky Inpainting Inference

scripts/mask_generator.py

@@ -4,10 +4,8 @@ from logger import rich_logger as l
 from ultralytics import YOLO
 import cv2
 from config import yolo_model
-
-
-
-
+from pathlib import Path
+import PIL.ImageOps
 
 
 
@@ -15,43 +13,31 @@ from config import yolo_model
 def generate_mask(image_path: str) -> np.ndarray:
     """Method to segment image
     Args:
-        image (Image): input image
+        image_path (str): path to input image
     Returns:
-        Image: segmented image
+        np.ndarray: segmented image mask
     """
-    model = YOLO(model=yolo_model)
-    results = model(image_path)
+    model = YOLO(model=yolo_model)  # Initialize YOLO model
+    results = model(image_path)  # Perform object detection
     for result in results:
         orig_img = result.orig_img
         masks = result.masks.xy
         height, width = result.orig_img.shape[:2]
-        background = np.ones((height, width, 3), dtype=np.uint8) * 255
+        mask_img = np.ones((height, width), dtype=np.uint8) * 255  # Initialize mask with white background
         
         for mask in masks:
-           mask = mask.astype(int)
-           mask_img = np.zeros_like(orig_img)
-           cv2.fillPoly(mask_img, [mask], (255, 255, 255))
-           mask_img = np.array(mask_img)
-           orig_img = np.array(orig_img)
-           
-    return mask_img, orig_img
+            mask = mask.astype(int)
+            cv2.fillPoly(mask_img, [mask], 0)  # Fill mask with detected object areas
+
+    return mask_img
 
-def invert_mask(mask_image: np.ndarray) -> np.ndarray:
+def invert_mask(mask_image: Image) -> np.ndarray:
     """Method to invert mask
     Args:
         mask_image (np.ndarray): input mask image
     Returns:
         np.ndarray: inverted mask image
     """
-    inverted_mask_image = cv2.bitwise_not(mask_image)
-    cv2.imwrite('invert_mask.jpg', inverted_mask_image)
+    inverted_mask_image =PIL.ImageOps.invert(mask_image)
     return inverted_mask_image
 
- 
-if __name__ == "__main__":
-    image = Image.open("../sample_data/example1.jpg")
-    mask_img,orig_image = generate_mask(image_path='../sample_data/example1.jpg')
-    invert_mask(mask_image=mask_img)
-    
-    
-

scripts/models.py

@@ -5,13 +5,13 @@ from clear_memory import clear_memory
 from typing import List
 import numpy as np
 import torch
-from PIL import Image
+from PIL import Image,ImageFilter,ImageOps
 from mask_generator import invert_mask
 from diffusers.utils import load_image
 from pipeline import fetch_control_pipeline,fetch_kandinsky_pipeline,fetch_kandinsky_prior_pipeline,fetch_kandinsky_img2img_pipeline
 from config import controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name
 import cv2 
-import PIL.ImageOps
+
 from transformers import pipeline
 
 
@@ -126,14 +126,12 @@ def kandinsky_inpainting_inference(prompt, negative_prompt, image, mask_image):
     Returns:
         PIL.Image.Image: The output inpainted image.
     """
-    pipe = fetch_kandinsky_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image)
-    output_image = pipe(prompt=prompt, negative_prompt=negative_prompt, image=image, mask_image=mask_image).images[0]
-    return output_image
-def kandinsky_inpainting_inference(prompt,negative_prompt,image,mask_image):
     pipe = fetch_kandinsky_pipeline(controlnet_adapter_model_name, controlnet_base_model_name,kandinsky_model_name, image)
-    output_image = pipe(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image).images[0]
+    output_image = pipe(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image,num_inference_steps=200,strength=1.0).images[0]
     return output_image
 
+    
+
 def kandinsky_controlnet_inpainting_inference(prompt, negative_prompt, image, hint, generator=torch.Generator(device="cuda").manual_seed(43)):
     """
     Perform inpainting inference using the Kandinsky ControlNet model.
@@ -150,28 +148,15 @@ def kandinsky_controlnet_inpainting_inference(prompt, negative_prompt, image, hi
 
     """
     prior_pipe = fetch_kandinsky_prior_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image)
-    img_embed = prior_pipe(prompt=prompt, image=image, strength=0.85, generator=generator)
+    img_embed = prior_pipe(prompt=prompt, image=image, strength=1.0, generator=generator)
     negative_embed = prior_pipe(prompt=negative_prompt, image=image, strength=1, generator=generator)
     controlnet_pipe = fetch_kandinsky_img2img_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image)
-    image = controlnet_pipe(image=image, strength=0.5, image_embeds=img_embed.image_embeds, negative_image_embeds=negative_embed.image_embeds, hint=hint, num_inference_steps=50, generator=generator, height=768, width=768).images[0]
+    image = controlnet_pipe(image=image, strength=1.0, image_embeds=img_embed.image_embeds, negative_image_embeds=negative_embed.image_embeds, hint=hint, num_inference_steps=200, generator=generator, height=768, width=768).images[0]
     return image
 
   
 
-if __name__ == '__main__':
-    l.info("Kandinsky Inpainting Inference")
-    image = load_image('/home/product_diffusion_api/sample_data/example2.jpg')
-    image = image.resize((768, 768))
-    mask_image = load_image('/home/product_diffusion_api/scripts/invert_mask.jpg')
-    mask_image = mask_image.resize((768,768))
-    prompt = "Product in a GYM 8k ultrarealistic "
-    negative_prompt="lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"
-    output_image = kandinsky_inpainting_inference(prompt,negative_prompt,image,mask_image)
-    output_image=output_image.resize((768,768))
-    depth_estimator = pipeline("depth-estimation")
-    hint = make_hint(output_image, depth_estimator).unsqueeze(0).half().to("cuda")
-    final_output_image = kandinsky_controlnet_inpainting_inference(prompt,negative_prompt,image, hint)
-    
+
     
     
     

scripts/run.py

@@ -0,0 +1,45 @@
+import argparse
+import os
+from mask_generator import generate_mask, invert_mask
+from models import kandinsky_inpainting_inference, load_image
+from PIL import Image
+
+def main(args):
+    # Generate mask
+    mask = generate_mask(args.image_path)
+    mask_image = Image.fromarray(mask)
+
+    # Save original mask
+    original_mask_path = os.path.join(args.mask_dir, f'{args.uid}_original_mask.jpg')
+    mask_image.save(original_mask_path)
+    
+    
+    # Invert mask
+    mask_image = load_image(original_mask_path)
+    inverted_mask = invert_mask(mask_image)
+    inverted_mask_path = os.path.join(args.mask_dir, f'{args.uid}_inverted_mask.jpg')
+    inverted_mask.save(inverted_mask_path)
+
+    # Load mask and image
+    invert_mask_image = load_image(inverted_mask_path)
+    image = load_image(args.image_path)
+
+    # Perform inpainting
+    output_image = kandinsky_inpainting_inference(args.prompt, args.negative_prompt, image, mask_image)
+
+    # Save output image
+    output_image_path = os.path.join(args.output_dir, f'{args.uid}_output.jpg')
+    output_image.save(output_image_path)
+    
+    
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Perform Kandinsky inpainting on an image.')
+    parser.add_argument('--image_path', type=str, required=True, help='Path to the input image.')
+    parser.add_argument('--prompt', type=str, required=True, help='Prompt for the Kandinsky inpainting.')
+    parser.add_argument('--negative_prompt', type=str, required=True, help='Negative prompt for the Kandinsky inpainting.')
+    parser.add_argument('--output_dir', type=str, required=True, help='Directory to save the output image.')
+    parser.add_argument('--mask_dir', type=str, required=True, help='Directory to save the mask image.')
+    parser.add_argument('--uid', type=str, required=True, help='Unique identifier for the image and mask.')
+
+    args = parser.parse_args()
+    main(args)