Spaces:
Running on Zero
Running on Zero
File size: 4,090 Bytes
e994c16 63c4f20 e994c16 2fb233c e994c16 63c4f20 e994c16 63c4f20 e994c16 | 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 136 137 138 139 | """Image preprocessing and visualization utilities."""
from __future__ import annotations
import hashlib
import io
import base64
from pathlib import Path
from typing import Any, Iterable
from PIL import Image, ImageDraw, ImageFont, ImageOps
from data.schemas import LABEL_DISPLAY_NAMES, bbox_to_pixels
LABEL_STYLE = {
"dust": ((245, 158, 11), 2),
"dirt": ((217, 119, 6), 2),
"scratch": ((220, 38, 38), 3),
"long_hair": ((124, 58, 237), 2),
"short_hair": ((8, 145, 178), 2),
"emulsion_damage": ((226, 232, 240), 3),
"chemical_stain": ((22, 163, 74), 3),
"light_leak": ((244, 114, 182), 3),
}
DEFAULT_STYLE = ((255, 255, 255), 2)
def load_image(image: str | Path | Image.Image) -> Image.Image:
"""Load an image-like value and return RGB PIL Image."""
if isinstance(image, Image.Image):
pil = image
else:
pil = Image.open(image)
pil = ImageOps.exif_transpose(pil)
if pil.mode == "RGBA":
background = Image.new("RGB", pil.size, (24, 22, 20))
background.paste(pil, mask=pil.getchannel("A"))
return background
if pil.mode != "RGB":
return pil.convert("RGB")
return pil.copy()
def image_to_png_bytes(image: Image.Image) -> bytes:
buf = io.BytesIO()
load_image(image).save(buf, format="PNG", optimize=True)
return buf.getvalue()
def image_to_data_uri(
image: Image.Image,
*,
max_side: int = 1800,
image_format: str = "JPEG",
quality: int = 92,
) -> str:
"""Return a browser-openable image data URI for review previews."""
pil = resize_for_preview(load_image(image), max_side=max_side)
fmt = image_format.upper()
buf = io.BytesIO()
if fmt in {"JPG", "JPEG"}:
pil = pil.convert("RGB")
pil.save(buf, format="JPEG", quality=quality, optimize=True)
mime = "image/jpeg"
elif fmt == "PNG":
pil.save(buf, format="PNG", optimize=True)
mime = "image/png"
else:
raise ValueError(f"unsupported image_format: {image_format}")
encoded = base64.b64encode(buf.getvalue()).decode("ascii")
return f"data:{mime};base64,{encoded}"
def image_sha256(image: Image.Image | bytes) -> str:
if isinstance(image, bytes):
payload = image
else:
payload = image_to_png_bytes(image)
return hashlib.sha256(payload).hexdigest()
def resize_for_preview(image: Image.Image, max_side: int = 1400) -> Image.Image:
pil = load_image(image)
if max(pil.size) <= max_side:
return pil
out = pil.copy()
out.thumbnail((max_side, max_side), Image.Resampling.LANCZOS)
return out
def draw_defects(
image: Image.Image,
defects: Iterable[dict[str, Any]],
*,
title: str | None = None,
max_boxes: int = 300,
) -> Image.Image:
"""Draw normalized defect boxes onto an RGB copy of an image."""
out = load_image(image)
draw = ImageDraw.Draw(out)
width, height = out.size
font = ImageFont.load_default()
if title:
draw.rectangle((0, 0, min(width, 440), 24), fill=(12, 10, 9))
draw.text((8, 6), title, fill=(254, 243, 199), font=font)
drawn = 0
for defect in defects:
if drawn >= max_boxes:
break
label = str(defect.get("label", "unknown"))
pixels = bbox_to_pixels(defect.get("bbox"), width, height)
if pixels is None:
continue
x_min, y_min, x_max, y_max = pixels
color, line_width = LABEL_STYLE.get(label, DEFAULT_STYLE)
draw.rectangle((x_min, y_min, x_max, y_max), outline=color, width=line_width)
label_text = LABEL_DISPLAY_NAMES.get(label, label)
text_bbox = draw.textbbox((x_min, max(0, y_min - 16)), label_text, font=font)
draw.rectangle(text_bbox, fill=(12, 10, 9))
draw.text((text_bbox[0] + 1, text_bbox[1]), label_text, fill=color, font=font)
drawn += 1
return out
__all__ = [
"DEFAULT_STYLE",
"LABEL_STYLE",
"draw_defects",
"image_sha256",
"image_to_data_uri",
"image_to_png_bytes",
"load_image",
"resize_for_preview",
]
|