Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| def find_bounding_box(mask): | |
| """Find the bounding box around the largest contour in the mask.""" | |
| mask = mask.astype(np.uint8) | |
| _, binary_mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY) | |
| contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| if contours: | |
| x, y, w, h = cv2.boundingRect(np.vstack(contours)) | |
| return x, y, w, h | |
| return None | |
| def overlay_logo_on_image(original_img, mask, logo): | |
| """Overlay a scaled logo while maintaining aspect ratio and centering it in the bounding box.""" | |
| bbox = find_bounding_box(mask) | |
| if bbox is None: | |
| print("No bounding box found. Returning original image.") | |
| return original_img | |
| x, y, w, h = bbox | |
| print(f"Bounding box found at: x={x}, y={y}, w={w}, h={h}") | |
| # Get original logo dimensions | |
| logo_h, logo_w = logo.shape[:2] | |
| # Compute the scaling factor while maintaining aspect ratio | |
| scale = min(w / logo_w, h / logo_h) | |
| new_w, new_h = int(logo_w * scale), int(logo_h * scale) | |
| # Resize the logo while maintaining aspect ratio | |
| logo_resized = cv2.resize(logo, (new_w, new_h), interpolation=cv2.INTER_AREA) | |
| # Ensure logo_resized has 4 channels (RGBA) | |
| if logo_resized.shape[2] == 3: | |
| logo_resized = cv2.cvtColor(logo_resized, cv2.COLOR_RGB2RGBA) | |
| # Create a blank transparent canvas (RGBA) | |
| logo_canvas = np.zeros((h, w, 4), dtype=np.uint8) | |
| # Compute centering offsets | |
| x_offset = (w - new_w) // 2 | |
| y_offset = (h - new_h) // 2 | |
| # Place the resized logo at the center of the bounding box | |
| logo_canvas[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = logo_resized | |
| # Convert original image to 4-channel if necessary | |
| if original_img.shape[2] == 3: | |
| original_img = cv2.cvtColor(original_img, cv2.COLOR_RGB2RGBA) | |
| # Blend the logo into the image using alpha blending | |
| alpha = logo_canvas[:, :, 3] / 255.0 | |
| for c in range(3): | |
| original_img[y:y+h, x:x+w, c] = ( | |
| (1 - alpha) * original_img[y:y+h, x:x+w, c] + alpha * logo_canvas[:, :, c] | |
| ) | |
| return original_img | |
| def process_images(original_img, mask, logo, progress=gr.Progress()): | |
| """Process images: overlay logo using uploaded mask.""" | |
| original_img = cv2.cvtColor(np.array(original_img), cv2.COLOR_RGB2BGR) | |
| mask = np.array(mask) | |
| print(f"Logo type before conversion: {type(logo)}, dtype: {getattr(logo, 'dtype', 'N/A')}") | |
| logo = np.array(logo) | |
| if logo.dtype != np.uint8: | |
| logo = logo.astype(np.uint8) | |
| logo = cv2.cvtColor(np.array(logo), cv2.COLOR_RGB2BGR) | |
| if len(mask.shape) == 3: | |
| mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY) | |
| mask = (mask > 0).astype(np.uint8) * 255 | |
| print(f"Mask shape: {mask.shape}, unique values: {np.unique(mask)}") | |
| result_img = overlay_logo_on_image(original_img, mask, logo) | |
| result_img = cv2.cvtColor(result_img, cv2.COLOR_BGRA2RGBA) if result_img.shape[2] == 4 else cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB) | |
| return result_img | |
| with gr.Blocks() as demo : | |
| gr.Markdown("## Inpaint Logo on Uploaded Mask") | |
| with gr.Row() : | |
| with gr.Column() : | |
| with gr.Row(): | |
| original_img = gr.Image(label="Upload Original Image", type="numpy") | |
| with gr.Row(): | |
| mask = gr.Image(label="Upload Mask (Binary)", type="numpy") | |
| logo = gr.Image(label="Upload Logo", type="numpy") | |
| btn = gr.Button("Submit") | |
| with gr.Column() : | |
| output = gr.Image(label="Output Image", format="png") | |
| btn.click(fn=process_images, inputs=[original_img, mask, logo], outputs=output) | |
| demo.launch(share=True, debug=True, show_error=True) | |