Spaces:
Build error
Build error
| import cv2 | |
| import numpy as np | |
| from PIL import Image, ImageDraw, ImageFont | |
| import torch | |
| import torchvision.transforms as T | |
| from torchvision.models import resnet50 | |
| from scipy.ndimage import gaussian_filter | |
| import gradio as gr | |
| class TextClothBlender: | |
| def __init__(self, font_path: str): | |
| self.font_path = font_path | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.model = resnet50(pretrained=True).to(self.device) | |
| self.model.eval() | |
| def preprocess_cloth_image(self, cloth_image_path: str): | |
| # Load the cloth image | |
| cloth_image = cv2.imread(cloth_image_path) | |
| gray_image = cv2.cvtColor(cloth_image, cv2.COLOR_BGR2GRAY) | |
| # Detect texture using edge detection | |
| edges = cv2.Canny(gray_image, 50, 150) | |
| # Extract features using ResNet | |
| preprocess = T.Compose([ | |
| T.ToPILImage(), | |
| T.Resize((224, 224)), | |
| T.ToTensor(), | |
| T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| ]) | |
| input_tensor = preprocess(cloth_image).unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| features = self.model(input_tensor) | |
| return cloth_image, edges, features | |
| def render_text(self, text: str, font_size: int, text_color: tuple): | |
| # Create a blank image for text | |
| font = ImageFont.truetype(self.font_path, font_size) | |
| text_size = font.getsize(text) | |
| text_image = Image.new('RGBA', text_size, (255, 255, 255, 0)) | |
| draw = ImageDraw.Draw(text_image) | |
| draw.text((0, 0), text, font=font, fill=text_color) | |
| return np.array(text_image) | |
| def apply_perspective_transform(self, text_image: np.ndarray, cloth_image_shape: tuple): | |
| # Define points for perspective transformation | |
| h, w, _ = cloth_image_shape | |
| src_points = np.float32([[0, 0], [text_image.shape[1], 0], [0, text_image.shape[0]], [text_image.shape[1], text_image.shape[0]]]) | |
| dst_points = np.float32([[50, 50], [w - 50, 30], [50, h - 100], [w - 50, h - 120]]) | |
| matrix = cv2.getPerspectiveTransform(src_points, dst_points) | |
| warped_text = cv2.warpPerspective(text_image, matrix, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0)) | |
| return warped_text | |
| def blend_text_with_cloth(self, cloth_image: np.ndarray, text_image: np.ndarray, edges: np.ndarray): | |
| # Convert cloth and text images to the same size | |
| h, w, _ = cloth_image.shape | |
| text_resized = cv2.resize(text_image, (w, h), interpolation=cv2.INTER_AREA) | |
| # Convert text to grayscale for masking | |
| text_gray = cv2.cvtColor(text_resized, cv2.COLOR_RGBA2GRAY) | |
| _, text_mask = cv2.threshold(text_gray, 1, 255, cv2.THRESH_BINARY) | |
| # Apply displacement mapping using the edges | |
| displace_map = gaussian_filter(edges, sigma=5) | |
| displaced_text = cv2.addWeighted(text_resized, 0.5, displace_map[..., None], 0.5, 0) | |
| # Blend text and cloth using overlay mode | |
| blended = cv2.addWeighted(cloth_image, 0.7, displaced_text[..., :3], 0.3, 0) | |
| return blended | |
| def refine_output(self, blended_image: np.ndarray): | |
| # Apply Gaussian blur for smooth edges | |
| refined = cv2.GaussianBlur(blended_image, (5, 5), 0) | |
| return refined | |
| def process(self, cloth_image_path: str, text: str, font_size: int, text_color: tuple): | |
| # Step 1: Preprocess the cloth image | |
| cloth_image, edges, _ = self.preprocess_cloth_image(cloth_image_path) | |
| # Step 2: Render the text | |
| text_image = self.render_text(text, font_size, text_color) | |
| # Step 3: Apply perspective transformation | |
| warped_text = self.apply_perspective_transform(text_image, cloth_image.shape) | |
| # Step 4: Blend text with the cloth | |
| blended_image = self.blend_text_with_cloth(cloth_image, warped_text, edges) | |
| # Step 5: Refine the output | |
| final_image = self.refine_output(blended_image) | |
| return final_image | |
| # Define the Gradio interface | |
| def blend_text_on_cloth(cloth_image, text, font_size, text_color): | |
| font_path = "path/to/font.ttf" # Ensure the font file exists in the deployed environment | |
| blender = TextClothBlender(font_path) | |
| text_color = tuple(map(int, text_color.strip('()').split(','))) # Convert string to tuple | |
| cloth_image_path = "temp_cloth_image.jpg" | |
| # Save the uploaded cloth image temporarily | |
| cv2.imwrite(cloth_image_path, cv2.cvtColor(np.array(cloth_image), cv2.COLOR_RGB2BGR)) | |
| # Process the image | |
| result = blender.process(cloth_image_path, text, int(font_size), text_color) | |
| return result[:, :, ::-1] # Convert BGR to RGB for display | |
| iface = gr.Interface( | |
| fn=blend_text_on_cloth, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Cloth Image"), | |
| gr.Textbox(label="Text to Blend"), | |
| gr.Slider(10, 100, step=1, label="Font Size"), | |
| gr.Textbox(label="Text Color (R,G,B,A)", placeholder="e.g., 255,0,0,255") | |
| ], | |
| outputs=gr.Image(type="numpy", label="Blended Output") | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() |