Merge pull request #30 from boostcampaitech5/T5165_이태순

T5165 이태순

GPUserver/model/inference_pipeline_with_Inpaint_Anything.py

@@ -1,8 +1,8 @@
 from model.img2img_model import StableDiffusionImg2Img
 from model.segment_anything_model import SAM
-from PIL import Image
+from PIL import Image,ImageChops
 import numpy as np 
-from model.utils import image_resize,image_segmentation,combine_image
+from model.utils import image_resize,image_segmentation,combine_image, check_size, mask_composit, mask_composit2
 from model.Inpaint_Anything.lama_inpaint import inpaint_img_with_lama
 from model.Inpaint_Anything.utils import save_array_to_img, dilate_mask
 
@@ -19,24 +19,54 @@ def __init__(self,check_point:str=check_point_dump_path):
     def load_lora(self,check_point:str):
         self.image_translation.load_Lora(check_point)
     def pipe(self,image:Image.Image,input_bbox:np.array,prompt:str,negative_prompt:str,inference_steps:int=100,strength:float=0.6)->Image.Image:
-        #1.입력 이미지에서 마스크 추출
+        #1. 입력 이미지에서 마스크 추출
         mask = self.segmentation.make_mask_with_bbox(image,input_bbox)
         """
         1번을 수행한다면 아래 2번과 3번 과정은 비동기로 작성해도 가능할 것 같다.
         추후 속도를 위해 개선할 예정!
         """
-        #2-1. 위에서 추출된 마스크 이외의 배경을 제거한 뒤, 이미지 변환
+        #2-1. 위에서 추출된 마스크 이외의 배경을 제거
         segment_img = image_segmentation(image,mask)
 
-        target_image = self.image_translation.inpaint(prompt=prompt,negative_prompt=negative_prompt,image=segment_img,num_inference_steps=inference_steps,strength=strength)
-
+
+        #2-2. 선택된 객체만 crop
+        crop_img = segment_img.crop((input_bbox[0],input_bbox[1],input_bbox[2],input_bbox[3]))
+
+        #2-3. 선택된 객체 resize
+        resized_crop_img = check_size(crop_img)
+
+        #3. Resize한 객체 캐릭터 생성
+        resized_target_image = self.image_translation.inpaint(prompt=prompt,negative_prompt=negative_prompt,image=resized_crop_img,num_inference_steps=inference_steps,strength=strength)
+
+        #3-1. 크기 Resize한 객체 원본 크기로 돌리기
+        x = input_bbox[2] - input_bbox[0]
+        y = input_bbox[3] - input_bbox[1]
+        target_image = resized_target_image.resize((x, y),Image.ANTIALIAS)
+
+        #4. 캐릭터 마스크 재생성
+        x, y = target_image.size
+        bbox = np.array([0,0,x,y])
+
+        crop_mask = self.segmentation.make_mask_with_bbox(target_image, bbox)
+        # SAM에서 마스크의 테두리 부분이 잘 Seg 되지 않는 현상 때문에 미세하게 crop 진행
+        crop_mask = crop_mask.crop((4,4,x-4,y-4))
+
+        # 마스크 반전
+        crop_mask = ImageChops.invert(crop_mask)
+
+        #5. 애니 생성 후 원본 이미지 크기의 마스크 만들기
+        after_mask = mask_composit(image=image,crop_mask=crop_mask, input_bbox=input_bbox)
+
+        #6. 위에서 추출된 마스크 이외의 배경 제거
+        after_segment_img = mask_composit2(image, target_image, input_bbox)
+
+        #7. 원본 배경만 남기기
         background = self.outpaint(image,mask,15)
-        background = image_resize(image=target_image,background=background)
 
-        target_mask = self.segment_image(image = target_image,input_bbox=input_bbox)
-        #4 이미지를 합치기 전에 target_image의 mask를 통해 배경과 객체를 분리
-        target_image,target_background = image_segmentation(image=target_image,mask=target_mask,background=background)  
-        #5. 4번 과정에서 얻은 결과를 통해 두이미지를 합친다.
+        #8. 이미지를 합치기 전에 target_image의 mask를 통해 배경과 객체를 분리
+        target_image,target_background = image_segmentation(image=after_segment_img,mask=after_mask,background=background)  
+
+        #9. 8번 과정에서 얻은 결과를 통해 두이미지를 합친다.
         result = combine_image(background=target_background,image=target_image)
         return result
 

GPUserver/model/utils.py

@@ -36,4 +36,29 @@ def combine_image(background:Image.Image,image:Image.Image) -> Image.Image:
     background,image = np.array(background), np.array(image)
     result_np = background+image
     result_pil = Image.fromarray(result_np)
-    return result_pil
+    return result_pil
+
+def check_size(img):
+    w,h = img.size
+    tmp = w if w > h else h
+    standard = 950
+
+    ratio = standard / tmp
+    img = img.resize((int(w * ratio), int(h * ratio)),Image.ANTIALIAS)
+    return img
+
+def mask_composit(image:Image.Image,crop_img:Image.Image,input_bbox:np.array)->Image.Image:
+    x,y = image.size
+    board = np.zeros((y,x),dtype=bool)
+    board[input_bbox[1]+4:input_bbox[3]-4, input_bbox[0]+4:input_bbox[2]-4] = np.array(crop_img)
+
+    board = Image.fromarray(board)
+    return board
+
+def mask_composit2(image:Image.Image,orig_img:Image.Image,input_bbox:np.array)->Image.Image:
+    x,y = image.size
+    board = np.zeros((y,x,3),dtype=np.uint8)
+    board[input_bbox[1]:input_bbox[3], input_bbox[0]:input_bbox[2]] = np.array(orig_img)
+    board = Image.fromarray(board)
+
+    return board