Spaces:
Paused
Paused
File size: 4,197 Bytes
c7159bc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """
Image Utilities
Helper functions for resize, restore, comparison, and mask visualization.
"""
from __future__ import annotations
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
def resize_for_processing(
image: Image.Image,
max_size: int = 1024,
) -> tuple[Image.Image, dict]:
"""
Resize image so the longest side ≤ max_size, maintaining aspect ratio.
Returns the resized image and info needed to restore the original size.
"""
w, h = image.size
scale = min(max_size / max(w, h), 1.0)
new_w = int(w * scale)
new_h = int(h * scale)
resized = image.resize((new_w, new_h), Image.LANCZOS)
return resized, {"original_w": w, "original_h": h, "scale": scale}
def restore_original_size(
processed: Image.Image,
original: Image.Image,
scale_info: dict,
) -> Image.Image:
"""Upscale processed image back to the original dimensions."""
ow, oh = scale_info["original_w"], scale_info["original_h"]
if processed.size == (ow, oh):
return processed
return processed.resize((ow, oh), Image.LANCZOS)
def create_comparison_image(
original: Image.Image,
ironed: Image.Image,
divider_width: int = 4,
label_height: int = 36,
) -> np.ndarray:
"""
Create a side-by-side before/after comparison image with labels.
"""
w, h = original.size
ironed_resized = ironed.resize((w, h), Image.LANCZOS)
canvas_w = w * 2 + divider_width
canvas_h = h + label_height
canvas = Image.new("RGB", (canvas_w, canvas_h), (240, 240, 240))
# paste images
canvas.paste(original, (0, label_height))
canvas.paste(ironed_resized, (w + divider_width, label_height))
# divider line
draw = ImageDraw.Draw(canvas)
draw.rectangle(
[w, 0, w + divider_width - 1, canvas_h],
fill=(255, 255, 255),
)
# labels
try:
font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 16)
except Exception:
font = ImageFont.load_default()
label_bg_before = (40, 40, 40)
label_bg_after = (29, 158, 117) # teal
# before label
draw.rectangle([0, 0, w, label_height - 1], fill=label_bg_before)
draw.text((12, 10), "BEFORE", fill=(255, 255, 255), font=font)
# after label
draw.rectangle([w + divider_width, 0, canvas_w, label_height - 1], fill=label_bg_after)
draw.text((w + divider_width + 12, 10), "AFTER (ironed)", fill=(255, 255, 255), font=font)
return np.array(canvas)
def overlay_wrinkle_mask(
image: Image.Image,
wrinkle_map: np.ndarray,
garment_mask: np.ndarray,
alpha: float = 0.55,
) -> np.ndarray:
"""
Overlay a heatmap of wrinkle intensity on the image.
Red = high wrinkle, yellow = medium, transparent = low.
"""
arr = np.array(image.convert("RGB"))
h, w = arr.shape[:2]
# Ensure map matches image size
if wrinkle_map.shape != (h, w):
wrinkle_map = cv2.resize(wrinkle_map, (w, h), interpolation=cv2.INTER_LINEAR)
if garment_mask.shape != (h, w):
garment_mask_rs = cv2.resize(
garment_mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST
).astype(bool)
else:
garment_mask_rs = garment_mask
# apply garment mask to wrinkle map
viz_map = wrinkle_map * garment_mask_rs.astype(float)
# normalize
viz_map = (viz_map / (viz_map.max() + 1e-6) * 255).astype(np.uint8)
# colourmap: COLORMAP_JET gives blue→green→yellow→red
heat = cv2.applyColorMap(viz_map, cv2.COLORMAP_JET)
heat = cv2.cvtColor(heat, cv2.COLOR_BGR2RGB)
# only show heatmap where wrinkle_map > threshold
threshold_mask = (viz_map > 40).astype(np.float32)[..., np.newaxis]
blended = (
heat.astype(float) * threshold_mask * alpha
+ arr.astype(float) * (1.0 - threshold_mask * alpha)
).clip(0, 255).astype(np.uint8)
# draw garment contour
contour_mask = garment_mask_rs.astype(np.uint8) * 255
contours, _ = cv2.findContours(contour_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(blended, contours, -1, (0, 255, 150), 2)
return blended
|