import os from pathlib import Path import cv2 import numpy as np from PIL import Image # ========================= # 설정 # ========================= BASE_DIR = Path(__file__).resolve().parent IMAGE_PATH = BASE_DIR / "input.png" MASK_PATH = BASE_DIR / "mask.png" OUTPUT_PATH = BASE_DIR / "result.png" OUTPUT_ORIGINAL_SIZE_PATH = BASE_DIR / "result_original_size.png" MASK_BINARY_OUTPUT = BASE_DIR / "mask_binary.png" EXTRACT_RGBA_OUTPUT = BASE_DIR / "extracted_cloth_rgba.png" EXTRACT_WHITE_OUTPUT = BASE_DIR / "extracted_cloth_white.png" EXTRACT_CROP_OUTPUT = BASE_DIR / "extracted_cloth_crop.png" COMPLETED_MASK_OUTPUT = BASE_DIR / "completed_cloth_mask.png" MISSING_MASK_OUTPUT = BASE_DIR / "missing_cloth_mask.png" COMPLETION_BASE_OUTPUT = BASE_DIR / "completion_base.png" COMPLETED_RGBA_OUTPUT = BASE_DIR / "completed_cloth_rgba.png" COMPLETED_WHITE_OUTPUT = BASE_DIR / "completed_cloth_white.png" COMPLETED_CROP_OUTPUT = BASE_DIR / "completed_cloth_crop.png" FRONT_MASK_OUTPUT = BASE_DIR / "front_cloth_mask.png" FRONT_OUTPUT = BASE_DIR / "front_cloth.png" FRONT_RGBA_OUTPUT = BASE_DIR / "front_cloth_rgba.png" FRONT_WHITE_OUTPUT = BASE_DIR / "front_cloth_white.png" FRONT_CROP_OUTPUT = BASE_DIR / "front_cloth_crop.png" INVERT_MASK = False COMPLETE_OCCLUDED_CLOTH = True FRONT_VIEW_NORMALIZE = True RUN_INPAINT = os.environ.get("RUN_INPAINT", "1") != "0" USE_DIFFUSION_REFINEMENT = os.environ.get("USE_DIFFUSION_REFINEMENT", "0") == "1" INPAINT_LONG_SIDE = 768 COMPLETE_LOWER_START_RATIO = 0.32 MISSING_MASK_DILATE_KERNEL = 23 MISSING_MASK_BLUR_KERNEL = 17 INPAINT_STRENGTH = 0.25 if COMPLETE_OCCLUDED_CLOTH else 0.85 TEXTURE_SOURCE_EDGE_MARGIN = 34 TEXTURE_SAFE_ERODE_KERNEL = 23 TEXTURE_REPAIR_DILATE_KERNEL = 43 TEXTURE_REPAIR_BLUR_KERNEL = 31 SEAMLESS_BLEND_PATCH = False FRONT_WIDTH_SCALE = 1.04 FRONT_BODY_VERTICAL_RATIO = 0.36 FRONT_BODY_HALF_WIDTH_PERCENTILE = 82 FRONT_MIN_ROW_WIDTH_RATIO = 0.08 def normalize_mask(mask, size): if mask.size != size: mask = mask.resize(size, Image.Resampling.NEAREST) mask_np = np.array(mask.convert("L")) mask_bin = mask_np > 127 if INVERT_MASK: mask_bin = ~mask_bin return (mask_bin.astype(np.uint8) * 255) def save_extracted_images( image, mask_np, rgba_output=EXTRACT_RGBA_OUTPUT, white_output=EXTRACT_WHITE_OUTPUT, crop_output=EXTRACT_CROP_OUTPUT, ): image_np = np.array(image) alpha_np = mask_np.astype(np.uint8) rgba_np = np.dstack([image_np, alpha_np]) extracted_rgba = Image.fromarray(rgba_np) extracted_rgba.save(rgba_output) white_bg = np.ones_like(image_np, dtype=np.uint8) * 255 extracted_white = np.where(alpha_np[..., None] > 127, image_np, white_bg) Image.fromarray(extracted_white.astype(np.uint8)).save(white_output) ys, xs = np.where(alpha_np > 127) if xs.size == 0 or ys.size == 0: print("마스크 영역이 없습니다.") return bbox = ( int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1, ) extracted_rgba.crop(bbox).save(crop_output) print("마스크 bbox:", bbox) def keep_largest_component(mask_np): mask_bin = (mask_np > 127).astype(np.uint8) num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_bin, connectivity=8) if num_labels <= 1: return mask_np largest_label = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA])) return ((labels == largest_label).astype(np.uint8) * 255) def create_completed_cloth_mask(mask_np): clean_mask = keep_largest_component(mask_np) clean_mask = cv2.morphologyEx( clean_mask, cv2.MORPH_CLOSE, np.ones((7, 7), np.uint8), iterations=1, ) ys, xs = np.where(clean_mask > 127) if xs.size == 0 or ys.size == 0: return clean_mask height, width = clean_mask.shape y_min = int(ys.min()) y_max = int(ys.max()) garment_height = max(y_max - y_min, 1) center_x = int(np.median(xs)) y_start = y_min + int(garment_height * COMPLETE_LOWER_START_RATIO) lower_mask = np.zeros_like(clean_mask) lower_mask[y_start:y_max + 1] = clean_mask[y_start:y_max + 1] lower_ys, lower_xs = np.where(lower_mask > 127) if lower_xs.size < 3: return clean_mask points = np.column_stack([lower_xs, lower_ys]).astype(np.int32) hull = cv2.convexHull(points) hull_mask = np.zeros_like(clean_mask) cv2.fillConvexPoly(hull_mask, hull, 255) yy, xx = np.indices((height, width)) right_lower_region = ( (yy >= y_start) & (yy <= y_max) & (xx >= center_x) ) completed_mask = np.where( (hull_mask > 127) & right_lower_region, 255, clean_mask, ).astype(np.uint8) upper_right_edges = [] upper_band_end = min(y_max, y_start + int(garment_height * 0.18)) for y in range(y_start, upper_band_end + 1): row_xs = np.where(clean_mask[y] > 127)[0] if row_xs.size >= 24: upper_right_edges.append(int(row_xs.max())) if upper_right_edges: upper_right = int(np.percentile(upper_right_edges, 90)) else: upper_right = int(xs.max()) lower_right = int(center_x + (upper_right - center_x) * 0.72) for y in range(y_start, y_max + 1): row_xs = np.where(clean_mask[y] > 127)[0] if row_xs.size < 8: continue t = (y - y_start) / max(y_max - y_start, 1) target_right = int(round((1.0 - t) * upper_right + t * lower_right)) target_right = min(width - 1, max(int(row_xs.max()), target_right)) completed_mask[y, center_x:target_right + 1] = 255 return cv2.morphologyEx( completed_mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8), iterations=1, ) def create_missing_mask(visible_mask_np, completed_mask_np): missing_mask = np.where( (completed_mask_np > 127) & (visible_mask_np <= 127), 255, 0, ).astype(np.uint8) if not missing_mask.any(): return missing_mask missing_mask = cv2.dilate( missing_mask, np.ones((MISSING_MASK_DILATE_KERNEL, MISSING_MASK_DILATE_KERNEL), np.uint8), iterations=1, ) missing_mask = np.where(completed_mask_np > 127, missing_mask, 0).astype(np.uint8) return cv2.GaussianBlur(missing_mask, (MISSING_MASK_BLUR_KERNEL, MISSING_MASK_BLUR_KERNEL), 0) def copy_row_texture_into_missing(base_np, image_np, visible_mask_np, completed_mask_np): hard_missing = (completed_mask_np > 127) & (visible_mask_np <= 127) repair_mask = cv2.dilate( hard_missing.astype(np.uint8) * 255, np.ones((TEXTURE_REPAIR_DILATE_KERNEL, TEXTURE_REPAIR_DILATE_KERNEL), np.uint8), iterations=1, ) > 127 repair_mask = repair_mask & (completed_mask_np > 127) original_base_np = base_np.copy() visible_safe = cv2.erode( (visible_mask_np > 127).astype(np.uint8) * 255, np.ones((TEXTURE_SAFE_ERODE_KERNEL, TEXTURE_SAFE_ERODE_KERNEL), np.uint8), iterations=1, ) > 127 for y in np.where(repair_mask.any(axis=1))[0]: missing_xs = np.where(repair_mask[y])[0] visible_xs = np.where(visible_safe[y])[0] if visible_xs.size == 0: visible_xs = np.where(visible_mask_np[y] > 127)[0] if visible_xs.size == 0: continue split_at = np.where(np.diff(missing_xs) > 1)[0] + 1 for run in np.split(missing_xs, split_at): if run.size == 0: continue source_xs = visible_xs[visible_xs < run[0] - TEXTURE_SOURCE_EDGE_MARGIN] if source_xs.size == 0: source_xs = visible_xs source_xs = source_xs[-run.size:] source_pixels = image_np[y, source_xs] if source_pixels.shape[0] < run.size: source_pixels = cv2.resize( source_pixels.reshape(1, source_pixels.shape[0], 3), (run.size, 1), interpolation=cv2.INTER_LINEAR, )[0] base_np[y, run] = source_pixels[:run.size] blend_alpha = cv2.GaussianBlur( repair_mask.astype(np.uint8) * 255, (TEXTURE_REPAIR_BLUR_KERNEL, TEXTURE_REPAIR_BLUR_KERNEL), 0, ) blend_alpha[hard_missing] = 255 alpha = (blend_alpha.astype(np.float32) / 255.0)[..., None] blended_np = original_base_np.astype(np.float32) * (1.0 - alpha) + base_np.astype(np.float32) * alpha return np.clip(blended_np, 0, 255).astype(np.uint8) def seamless_blend_patch(result_img, base_img, blend_mask_np): if not SEAMLESS_BLEND_PATCH: return result_img clone_mask = np.where(blend_mask_np > 10, 255, 0).astype(np.uint8) if not clone_mask.any(): return result_img clone_mask = cv2.dilate(clone_mask, np.ones((9, 9), np.uint8), iterations=1) ys, xs = np.where(clone_mask > 0) x1 = max(0, int(xs.min()) - 8) x2 = min(clone_mask.shape[1], int(xs.max()) + 9) y1 = max(0, int(ys.min()) - 8) y2 = min(clone_mask.shape[0], int(ys.max()) + 9) if x2 <= x1 or y2 <= y1: return result_img source_bgr = cv2.cvtColor(np.array(result_img), cv2.COLOR_RGB2BGR) dest_bgr = cv2.cvtColor(np.array(base_img), cv2.COLOR_RGB2BGR) source_roi = source_bgr[y1:y2, x1:x2] mask_roi = clone_mask[y1:y2, x1:x2] center = ((x1 + x2) // 2, (y1 + y2) // 2) try: blended_bgr = cv2.seamlessClone( source_roi, dest_bgr, mask_roi, center, cv2.MIXED_CLONE, ) except cv2.error: return result_img return Image.fromarray(cv2.cvtColor(blended_bgr, cv2.COLOR_BGR2RGB)) def smooth_vector(values, kernel_size): kernel_size = int(kernel_size) if kernel_size % 2 == 0: kernel_size += 1 if kernel_size < 3: return values return cv2.GaussianBlur( values.astype(np.float32).reshape(-1, 1), (1, kernel_size), 0, ).reshape(-1) def create_front_view(image, mask_np): mask_bin = mask_np > 127 ys, xs = np.where(mask_bin) if xs.size == 0 or ys.size == 0: return image, mask_np image_np = np.array(image) height, width = mask_np.shape y_min = int(ys.min()) y_max = int(ys.max()) garment_height = max(y_max - y_min, 1) lefts = np.full(height, np.nan, dtype=np.float32) rights = np.full(height, np.nan, dtype=np.float32) centers = np.full(height, np.nan, dtype=np.float32) for y in range(y_min, y_max + 1): row_xs = np.where(mask_bin[y])[0] if row_xs.size < 2: continue lefts[y] = float(row_xs.min()) rights[y] = float(row_xs.max()) centers[y] = (lefts[y] + rights[y]) / 2.0 row_widths = rights - lefts + 1 max_row_width = float(np.nanmax(row_widths)) min_row_width = max(2.0, max_row_width * FRONT_MIN_ROW_WIDTH_RATIO) valid_rows = np.where((~np.isnan(lefts)) & (row_widths >= min_row_width))[0] if valid_rows.size < 2: return image, mask_np y_min = int(valid_rows.min()) y_max = int(valid_rows.max()) garment_height = max(y_max - y_min, 1) target_center = int(round(np.nanmedian(centers[valid_rows]))) row_indices = np.arange(height) half_widths = np.maximum(target_center - lefts, rights - target_center) half_widths = np.interp( row_indices, valid_rows, half_widths[valid_rows], ) smooth_kernel = max(9, int(round(garment_height * 0.10))) half_widths = smooth_vector(half_widths, smooth_kernel) body_start = y_min + int(garment_height * FRONT_BODY_VERTICAL_RATIO) body_rows = valid_rows[valid_rows >= body_start] if body_rows.size: body_half = float(np.percentile(half_widths[body_rows], FRONT_BODY_HALF_WIDTH_PERCENTILE)) for y in body_rows: blend = min(1.0, max(0.0, (y - body_start) / max(garment_height * 0.18, 1.0))) half_widths[y] = (1.0 - blend) * half_widths[y] + blend * body_half front_np = np.ones_like(image_np, dtype=np.uint8) * 255 front_mask_np = np.zeros_like(mask_np, dtype=np.uint8) for y in valid_rows: source_xs = np.where(mask_bin[y])[0] if source_xs.size < 2: continue source_x1 = int(source_xs.min()) source_x2 = int(source_xs.max()) target_half = int(round(max(2.0, half_widths[y] * FRONT_WIDTH_SCALE))) target_x1 = max(0, target_center - target_half) target_x2 = min(width - 1, target_center + target_half) target_width = target_x2 - target_x1 + 1 if target_width <= 1: continue source_strip = image_np[y:y + 1, source_x1:source_x2 + 1] front_np[y, target_x1:target_x2 + 1] = cv2.resize( source_strip, (target_width, 1), interpolation=cv2.INTER_LINEAR, )[0] front_mask_np[y, target_x1:target_x2 + 1] = 255 front_mask_np = cv2.morphologyEx( front_mask_np, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8), iterations=1, ) front_np = np.where(front_mask_np[..., None] > 127, front_np, 255).astype(np.uint8) return Image.fromarray(front_np), front_mask_np def create_completion_base(image, visible_mask_np, completed_mask_np=None): image_np = np.array(image) white_bg = np.ones_like(image_np, dtype=np.uint8) * 255 base_np = np.where(visible_mask_np[..., None] > 127, image_np, white_bg) if completed_mask_np is not None: base_np = copy_row_texture_into_missing(base_np, image_np, visible_mask_np, completed_mask_np) return Image.fromarray(base_np.astype(np.uint8)) def resize_for_inpaint(image, mask): width, height = image.size long_side = max(width, height) scale = INPAINT_LONG_SIDE / long_side new_width = max(8, int(round((width * scale) / 8)) * 8) new_height = max(8, int(round((height * scale) / 8)) * 8) size = (new_width, new_height) if size == image.size: return image, mask return ( image.resize(size, Image.Resampling.LANCZOS), mask.resize(size, Image.Resampling.NEAREST), ) # ========================= # 이미지 / 마스크 로드 # ========================= image = Image.open(IMAGE_PATH).convert("RGB") mask = Image.open(MASK_PATH).convert("L") # 마스크는 부드러운 리사이즈를 쓰면 경계가 회색으로 번져서 영역이 틀어진다. mask_np = normalize_mask(mask, image.size) mask = Image.fromarray(mask_np) mask.save(MASK_BINARY_OUTPUT) save_extracted_images(image, mask_np) if COMPLETE_OCCLUDED_CLOTH: completed_mask_np = create_completed_cloth_mask(mask_np) missing_mask_np = create_missing_mask(mask_np, completed_mask_np) inpaint_base = create_completion_base(image, mask_np, completed_mask_np) inpaint_mask = Image.fromarray(missing_mask_np) Image.fromarray(completed_mask_np).save(COMPLETED_MASK_OUTPUT) Image.fromarray(missing_mask_np).save(MISSING_MASK_OUTPUT) inpaint_base.save(COMPLETION_BASE_OUTPUT) save_extracted_images( inpaint_base, completed_mask_np, COMPLETED_RGBA_OUTPUT, COMPLETED_WHITE_OUTPUT, COMPLETED_CROP_OUTPUT, ) else: completed_mask_np = mask_np missing_mask_np = mask_np inpaint_base = image inpaint_mask = mask if FRONT_VIEW_NORMALIZE: front_image, front_mask_np = create_front_view(inpaint_base, completed_mask_np) front_image.save(FRONT_OUTPUT) Image.fromarray(front_mask_np).save(FRONT_MASK_OUTPUT) save_extracted_images( front_image, front_mask_np, FRONT_RGBA_OUTPUT, FRONT_WHITE_OUTPUT, FRONT_CROP_OUTPUT, ) inpaint_base = front_image completed_mask_np = front_mask_np missing_mask_np = np.zeros_like(front_mask_np, dtype=np.uint8) inpaint_mask = Image.fromarray(missing_mask_np) if not RUN_INPAINT or not USE_DIFFUSION_REFINEMENT: preview_image, _ = resize_for_inpaint(inpaint_base, inpaint_mask) preview_image.save(OUTPUT_PATH) inpaint_base.save(OUTPUT_ORIGINAL_SIZE_PATH) save_extracted_images( inpaint_base, completed_mask_np, COMPLETED_RGBA_OUTPUT, COMPLETED_WHITE_OUTPUT, COMPLETED_CROP_OUTPUT, ) print("완료:", OUTPUT_PATH) print("원본 크기 결과:", OUTPUT_ORIGINAL_SIZE_PATH) print("완성 마스크:", COMPLETED_MASK_OUTPUT) print("복원할 영역 마스크:", MISSING_MASK_OUTPUT) print("정면화 결과:", FRONT_OUTPUT) print("완성 RGBA:", COMPLETED_RGBA_OUTPUT) raise SystemExit import torch from diffusers import StableDiffusionInpaintPipeline DEVICE = "cuda" if torch.cuda.is_available() else "cpu" inpaint_mask_original_size = inpaint_mask inpaint_image, inpaint_mask = resize_for_inpaint(inpaint_base, inpaint_mask_original_size) # ========================= # Inpainting 모델 로드 # ========================= pipe = StableDiffusionInpaintPipeline.from_pretrained( "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32 ) pipe = pipe.to(DEVICE) # 메모리 절약 옵션 pipe.enable_attention_slicing() # ========================= # 프롬프트 # ========================= prompt = """ complete only the missing occluded part of the black and white striped t-shirt, continue the horizontal stripes, same fabric texture and natural folds, same lighting, photorealistic """ negative_prompt = """ deformed, distorted, bad anatomy, extra limbs, blurry, low quality, cartoon, painting, unrealistic, skin, arm, hand, fingers, logo, text """ # ========================= # Inpainting 실행 # ========================= result = pipe( prompt=prompt, negative_prompt=negative_prompt, image=inpaint_image, mask_image=inpaint_mask, width=inpaint_image.width, height=inpaint_image.height, strength=INPAINT_STRENGTH, guidance_scale=7.5, num_inference_steps=40 ).images[0] # ========================= # 저장 # ========================= result = Image.composite(result, inpaint_image, inpaint_mask) result.save(OUTPUT_PATH) result_original_size = result.resize(image.size, Image.Resampling.LANCZOS) result_original_size = Image.composite(result_original_size, inpaint_base, inpaint_mask_original_size) if COMPLETE_OCCLUDED_CLOTH: result_original_size = seamless_blend_patch(result_original_size, inpaint_base, missing_mask_np) result_original_size.save(OUTPUT_ORIGINAL_SIZE_PATH) save_extracted_images( result_original_size, completed_mask_np, COMPLETED_RGBA_OUTPUT, COMPLETED_WHITE_OUTPUT, COMPLETED_CROP_OUTPUT, ) print("완료:", OUTPUT_PATH) print("원본 크기 결과:", OUTPUT_ORIGINAL_SIZE_PATH) print("RGBA 추출:", EXTRACT_RGBA_OUTPUT) print("흰 배경 추출:", EXTRACT_WHITE_OUTPUT) print("crop 추출:", EXTRACT_CROP_OUTPUT) print("완성 마스크:", COMPLETED_MASK_OUTPUT) print("복원할 영역 마스크:", MISSING_MASK_OUTPUT) print("정면화 결과:", FRONT_OUTPUT) print("완성 RGBA:", COMPLETED_RGBA_OUTPUT)