import cv2 import torch import numpy as np from PIL import Image from diffusers import StableDiffusionInpaintPipeline # -------------------------------------------------- # 설정 # -------------------------------------------------- IMAGE_PATH = "image.jpg" CLOTH_MASK_PATH = "mask.png" OUTPUT_EXTRACTED = "extracted_cloth.png" OUTPUT_RESULT = "product_cloth.png" OUTPUT_INPAINT_MASK = "inpaint_mask.png" OUTPUT_INPAINT_BASE = "inpaint_base.png" OUTPUT_PRODUCT_MASK = "product_shape_mask.png" OUTPUT_BRAND_MASK = "brand_removal_mask.png" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" MASK_DILATE_KERNEL = 3 MASK_BLUR_KERNEL = 5 NUM_INFERENCE_STEPS = 45 GUIDANCE_SCALE = 8.0 INPAINT_STRENGTH = 0.80 GENERATOR_SEED = 42 STRAIGHTEN_CLOTH = True CENTER_CLOTH = True MAX_STRAIGHTEN_ANGLE = 22.0 PRODUCT_SILHOUETTE = True REMOVE_LOWER_LEFT_BRAND = True ATTENTION_POSE = False ATTENTION_SLEEVE_WIDTH_RATIO = 0.34 ATTENTION_SLEEVE_OVERLAP_RATIO = 0.22 # True로 두면 옷 주변 배경까지 재생성되어 옷걸이/지지대가 생길 수 있다. POSTURE_REDRAW_MASK = False POSTURE_MASK_PADDING_X = 0.16 POSTURE_MASK_PADDING_TOP = 0.08 POSTURE_MASK_PADDING_BOTTOM = 0.20 def estimate_cloth_angle(mask_np): """마스크의 세로 중심선을 기준으로 옷이 기울어진 각도를 추정한다.""" mask_bin = mask_np > 127 rows = np.where(mask_bin.any(axis=1))[0] if rows.size < 12: return 0.0 y_min = int(rows[0]) y_max = int(rows[-1]) garment_height = max(y_max - y_min, 1) ys = [] centers = [] for y in range(y_min, y_max + 1): rel_y = (y - y_min) / garment_height if rel_y < 0.08 or rel_y > 0.92: continue xs = np.where(mask_bin[y])[0] if xs.size < 3: continue ys.append(y) centers.append(np.median(xs)) if len(ys) < 12: return 0.0 slope, _ = np.polyfit(np.asarray(ys, dtype=np.float32), np.asarray(centers, dtype=np.float32), 1) angle = float(np.degrees(np.arctan(slope))) return float(np.clip(angle, -MAX_STRAIGHTEN_ANGLE, MAX_STRAIGHTEN_ANGLE)) def warp_image_and_mask(image, mask, matrix, size): image_np = np.array(image) mask_np = np.array(mask) warped_image = cv2.warpAffine( image_np, matrix, size, flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=(255, 255, 255) ) warped_mask = cv2.warpAffine( mask_np, matrix, size, flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=0 ) return Image.fromarray(warped_image), Image.fromarray(warped_mask) def straighten_cloth(image, mask): mask_np = np.array(mask) angle = estimate_cloth_angle(mask_np) if abs(angle) < 0.5: return image, mask, angle width, height = image.size center = (width / 2, height / 2) matrix = cv2.getRotationMatrix2D(center, -angle, 1.0) image, mask = warp_image_and_mask(image, mask, matrix, (width, height)) return image, mask, angle def center_cloth_on_canvas(image, mask): mask_np = np.array(mask) ys, xs = np.where(mask_np > 127) if xs.size == 0 or ys.size == 0: return image, mask, (0.0, 0.0) width, height = image.size cloth_center_x = (float(xs.min()) + float(xs.max())) / 2 cloth_center_y = (float(ys.min()) + float(ys.max())) / 2 tx = (width / 2) - cloth_center_x ty = (height / 2) - cloth_center_y matrix = np.float32([[1, 0, tx], [0, 1, ty]]) image, mask = warp_image_and_mask(image, mask, matrix, (width, height)) return image, mask, (tx, ty) def create_posture_redraw_mask(mask_np): """소매와 옷 실루엣을 정자세로 다시 그릴 수 있도록 의류 주변 박스까지 마스크에 포함한다.""" ys, xs = np.where(mask_np > 127) if xs.size == 0 or ys.size == 0: return mask_np height, width = mask_np.shape x_min = int(xs.min()) x_max = int(xs.max()) y_min = int(ys.min()) y_max = int(ys.max()) cloth_width = max(x_max - x_min, 1) cloth_height = max(y_max - y_min, 1) pad_x = int(max(24, cloth_width * POSTURE_MASK_PADDING_X)) pad_top = int(max(18, cloth_height * POSTURE_MASK_PADDING_TOP)) pad_bottom = int(max(36, cloth_height * POSTURE_MASK_PADDING_BOTTOM)) x1 = max(0, x_min - pad_x) x2 = min(width - 1, x_max + pad_x) y1 = max(0, y_min - pad_top) y2 = min(height - 1, y_max + pad_bottom) redraw_mask = np.zeros_like(mask_np, dtype=np.uint8) redraw_mask[y1:y2 + 1, x1:x2 + 1] = 255 return redraw_mask def create_attention_pose_redraw_mask(mask_np): """벌어진 소매를 지우고 몸통 옆 아래 방향 소매를 새로 그릴 공간을 만든다.""" mask_bin = mask_np > 127 ys, xs = np.where(mask_bin) if xs.size == 0 or ys.size == 0: return mask_np height, width = mask_np.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)) lower_y1 = y_min + int(garment_height * 0.50) lower_y2 = y_min + int(garment_height * 0.86) torso_bounds = [] for y in range(lower_y1, lower_y2 + 1): row_xs = np.where(mask_bin[y])[0] if row_xs.size < 8: continue split_at = np.where(np.diff(row_xs) > 1)[0] + 1 runs = np.split(row_xs, split_at) central_run = min( runs, key=lambda run: 0 if run[0] <= center_x <= run[-1] else min(abs(run[0] - center_x), abs(run[-1] - center_x)) ) if central_run.size >= 24: torso_bounds.append((int(central_run[0]), int(central_run[-1]))) if torso_bounds: torso_x1 = int(np.median([bounds[0] for bounds in torso_bounds])) torso_x2 = int(np.median([bounds[1] for bounds in torso_bounds])) else: torso_x1 = int(np.percentile(xs, 35)) torso_x2 = int(np.percentile(xs, 65)) torso_width = max(torso_x2 - torso_x1, 1) sleeve_width = int(np.clip(torso_width * ATTENTION_SLEEVE_WIDTH_RATIO, 34, 58)) overlap = int(max(12, sleeve_width * ATTENTION_SLEEVE_OVERLAP_RATIO)) sleeve_y1 = y_min + int(garment_height * 0.18) sleeve_y2 = min(height - 1, y_min + int(garment_height * 0.98)) redraw_mask = np.where(mask_bin, 255, 0).astype(np.uint8) left_sleeve = np.array([ [max(0, torso_x1 - sleeve_width), sleeve_y1], [min(width - 1, torso_x1 + overlap), sleeve_y1], [min(width - 1, torso_x1 + overlap), sleeve_y2], [max(0, torso_x1 - int(sleeve_width * 0.78)), sleeve_y2], ], dtype=np.int32) right_sleeve = np.array([ [max(0, torso_x2 - overlap), sleeve_y1], [min(width - 1, torso_x2 + sleeve_width), sleeve_y1], [min(width - 1, torso_x2 + int(sleeve_width * 0.78)), sleeve_y2], [max(0, torso_x2 - overlap), sleeve_y2], ], dtype=np.int32) cv2.fillPoly(redraw_mask, [left_sleeve, right_sleeve], 255) return redraw_mask def create_product_sweatshirt_mask(size): """참고 이미지처럼 자연스러운 상품컷 스웨트셔츠 실루엣을 만든다.""" width, height = size scale_x = width / 512.0 scale_y = height / 512.0 def points(coords): return np.array( [[int(round(x * scale_x)), int(round(y * scale_y))] for x, y in coords], dtype=np.int32 ) mask_np = np.zeros((height, width), dtype=np.uint8) sweatshirt_outline = points([ (205, 108), (236, 121), (276, 121), (307, 108), (354, 128), (396, 160), (442, 306), (484, 456), (432, 476), (384, 396), (348, 440), (164, 440), (128, 396), (80, 476), (28, 456), (70, 306), (116, 160), (158, 128) ]) cv2.fillPoly(mask_np, [sweatshirt_outline], 255) # 안쪽 목 구멍은 흰 배경으로 남겨 두면 자연스러운 크루넥이 더 잘 나온다. cv2.ellipse( mask_np, (int(width * 0.5), int(height * 0.245)), (int(width * 0.105), int(height * 0.038)), 0, 0, 360, 0, -1 ) return mask_np def estimate_cloth_color(image_np, mask_np): cloth_pixels = image_np[mask_np > 127] if cloth_pixels.size == 0: return np.array([188, 190, 192], dtype=np.uint8) color = np.median(cloth_pixels, axis=0) color = np.clip(color, 172, 214) return color.astype(np.uint8) def create_product_inpaint_base(image_np, source_mask_np, target_mask_np): base_np = np.ones_like(image_np, dtype=np.uint8) * 255 cloth_color = estimate_cloth_color(image_np, source_mask_np) height, width = target_mask_np.shape yy, xx = np.indices((height, width)) x_norm = xx / max(width - 1, 1) y_norm = yy / max(height - 1, 1) vertical_fade = np.clip((y_norm - 0.18) / 0.62, 0.0, 1.0) fold_shadows = ( -7.0 * np.exp(-((x_norm - 0.43) / 0.028) ** 2) + 4.0 * np.exp(-((x_norm - 0.49) / 0.035) ** 2) + -8.0 * np.exp(-((x_norm - 0.57) / 0.030) ** 2) ) * vertical_fade side_shading = -10.0 * np.abs(x_norm - 0.5) rng = np.random.default_rng(GENERATOR_SEED) knit_noise = rng.normal(0.0, 2.2, (height, width, 1)) texture = fold_shadows[..., None] + side_shading[..., None] + knit_noise product_np = np.clip(cloth_color.astype(np.float32) + texture, 0, 255).astype(np.uint8) base_np[target_mask_np > 127] = product_np[target_mask_np > 127] return Image.fromarray(base_np) def remove_lower_left_brand_artifacts(result_img, garment_mask_np): """왼쪽 아래 몸판에 생성되는 작은 상표/패치만 감지해 주변 원단으로 메운다.""" result_np = np.array(result_img) height, width = garment_mask_np.shape x1 = int(width * 0.30) x2 = int(width * 0.45) y1 = int(height * 0.54) y2 = int(height * 0.76) roi = result_np[y1:y2, x1:x2] roi_garment = garment_mask_np[y1:y2, x1:x2] > 127 if roi.size == 0 or not roi_garment.any(): return result_img safe_garment = cv2.erode( roi_garment.astype(np.uint8) * 255, np.ones((17, 17), np.uint8), iterations=1 ) > 127 gray = cv2.cvtColor(roi, cv2.COLOR_RGB2GRAY) hsv = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV) garment_gray = gray[roi_garment] local_median = float(np.median(garment_gray)) # 회색 원단에서 튀는 작은 검은/흰 로고, 컬러 라벨, 텍스트 조각을 후보로 잡는다. candidate = ( safe_garment & ( (np.abs(gray.astype(np.float32) - local_median) > 34) | ((hsv[..., 1] > 34) & (np.abs(gray.astype(np.float32) - local_median) > 18)) ) ).astype(np.uint8) * 255 candidate = cv2.morphologyEx(candidate, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8)) num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(candidate, connectivity=8) removal_roi = np.zeros_like(candidate) for label in range(1, num_labels): area = int(stats[label, cv2.CC_STAT_AREA]) comp_x = int(stats[label, cv2.CC_STAT_LEFT]) comp_y = int(stats[label, cv2.CC_STAT_TOP]) comp_w = int(stats[label, cv2.CC_STAT_WIDTH]) comp_h = int(stats[label, cv2.CC_STAT_HEIGHT]) aspect = comp_w / max(comp_h, 1) is_logo_sized = 8 <= area <= 900 and 3 <= comp_w <= 62 and 3 <= comp_h <= 42 is_not_long_seam = aspect < 5.5 if is_logo_sized and is_not_long_seam: removal_roi[labels == label] = 255 if not removal_roi.any(): Image.fromarray(np.zeros((height, width), dtype=np.uint8)).save(OUTPUT_BRAND_MASK) return result_img removal_roi = cv2.dilate(removal_roi, np.ones((7, 7), np.uint8), iterations=1) removal_roi = cv2.GaussianBlur(removal_roi, (5, 5), 0) removal_mask = np.zeros((height, width), dtype=np.uint8) removal_mask[y1:y2, x1:x2] = removal_roi removal_mask = np.where(garment_mask_np > 127, removal_mask, 0).astype(np.uint8) Image.fromarray(removal_mask).save(OUTPUT_BRAND_MASK) cleaned_np = cv2.inpaint(result_np, removal_mask, 5, cv2.INPAINT_TELEA) return Image.fromarray(cleaned_np) # -------------------------------------------------- # 1. 이미지 로드 # -------------------------------------------------- image = Image.open(IMAGE_PATH).convert("RGB") image = image.resize((512, 512)) mask = Image.open(CLOTH_MASK_PATH).convert("L") mask = mask.resize((512, 512)) if STRAIGHTEN_CLOTH: image, mask, straightened_angle = straighten_cloth(image, mask) print(f"정자세 회전 보정: {-straightened_angle:.2f}도") if CENTER_CLOTH: image, mask, cloth_offset = center_cloth_on_canvas(image, mask) print(f"중앙 정렬 보정: x={cloth_offset[0]:.1f}, y={cloth_offset[1]:.1f}") # -------------------------------------------------- # 2. 마스크 확장 # 옷 형태를 자연스럽게 복원하기 위해 약간 넓힘 # -------------------------------------------------- mask_np = np.array(mask) if PRODUCT_SILHOUETTE: inpaint_mask_np = create_product_sweatshirt_mask(image.size) Image.fromarray(inpaint_mask_np).save(OUTPUT_PRODUCT_MASK) elif ATTENTION_POSE: inpaint_mask_np = create_attention_pose_redraw_mask(mask_np) elif POSTURE_REDRAW_MASK: inpaint_mask_np = create_posture_redraw_mask(mask_np) else: inpaint_mask_np = mask_np kernel = np.ones((MASK_DILATE_KERNEL, MASK_DILATE_KERNEL), np.uint8) expanded_mask = cv2.dilate(inpaint_mask_np, kernel, iterations=1) expanded_mask = cv2.GaussianBlur(expanded_mask, (MASK_BLUR_KERNEL, MASK_BLUR_KERNEL), 0) expanded_mask_img = Image.fromarray(expanded_mask) expanded_mask_img.save(OUTPUT_INPAINT_MASK) # -------------------------------------------------- # 3. 옷만 추출해서 흰 배경에 합성 # -------------------------------------------------- image_np = np.array(image) mask_3ch = np.stack([mask_np] * 3, axis=-1) white_bg = np.ones_like(image_np) * 255 cloth_only = np.where(mask_3ch > 127, image_np, white_bg) cloth_only_img = Image.fromarray(cloth_only.astype(np.uint8)) cloth_only_img.save(OUTPUT_EXTRACTED) if PRODUCT_SILHOUETTE: inpaint_base_img = create_product_inpaint_base(image_np, mask_np, inpaint_mask_np) inpaint_base_img.save(OUTPUT_INPAINT_BASE) else: inpaint_base_img = cloth_only_img # -------------------------------------------------- # 4. Inpainting 모델 로드 # -------------------------------------------------- pipe = StableDiffusionInpaintPipeline.from_pretrained( "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32 ) pipe = pipe.to(DEVICE) # -------------------------------------------------- # 5. 상품컷 스타일로 보정 # -------------------------------------------------- prompt = """ ecommerce photo of a plain unbranded light gray crewneck sweatshirt, empty garment, no hanger, no mannequin, no visible label, front view, centered, dropped shoulders, sleeves gently angled downward, ribbed collar, ribbed cuffs, ribbed bottom hem, heather gray cotton fleece texture, subtle vertical wrinkles, pure white background, soft studio lighting, sharp focus """ negative_prompt = """ hanger, hook, rack, mannequin, human, person, face, hands, skin, arms, model, side view, folded, flat lay, distorted, warped, extra sleeves, cropped, logo, label, tag, neck label, text, brand patch, black patch, low quality, blur """ generator = torch.Generator(device=DEVICE).manual_seed(GENERATOR_SEED) result = pipe( prompt=prompt, negative_prompt=negative_prompt, image=inpaint_base_img, mask_image=expanded_mask_img, num_inference_steps=NUM_INFERENCE_STEPS, guidance_scale=GUIDANCE_SCALE, strength=INPAINT_STRENGTH, generator=generator ).images[0] if REMOVE_LOWER_LEFT_BRAND: result = remove_lower_left_brand_artifacts(result, inpaint_mask_np) result.save(OUTPUT_RESULT) print("완료:", OUTPUT_RESULT)