Spaces:
Sleeping
Sleeping
File size: 4,023 Bytes
ad83752 | 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 | """
utils.py
--------
Image I/O and visualization helpers:
* safe image loading / validation
* colored prediction mask (dark gray / bright yellow)
* semi-transparent magenta crop overlay
* combined three-panel image + PNG export for download
"""
from __future__ import annotations
import io
import numpy as np
from PIL import Image, ImageDraw, ImageFont, UnidentifiedImageError
# --------------------------------------------------------------------------- #
# Visualization palette
# --------------------------------------------------------------------------- #
COLOR_BACKGROUND = (50, 50, 50) # dark gray -> Class 0 (No Crop)
COLOR_CROP = (255, 255, 0) # bright yellow -> Class 1 (Crop)
OVERLAY_COLOR = (255, 0, 255) # magenta overlay for crop regions
OVERLAY_ALPHA = 0.5 # overlay transparency (0..1)
class ImageLoadError(Exception):
"""Raised when an uploaded image is invalid or corrupted."""
def load_image(file) -> Image.Image:
"""
Safely load an uploaded file into an RGB PIL image.
Raises ImageLoadError on invalid / corrupted input.
"""
try:
image = Image.open(file)
image.load() # force a full decode so corrupted files fail here
except (UnidentifiedImageError, OSError, ValueError) as exc:
raise ImageLoadError(
"The uploaded file is not a valid image or is corrupted."
) from exc
return image.convert("RGB")
def create_color_mask(mask: np.ndarray) -> Image.Image:
"""Convert a {0, 1} mask into an RGB image (dark gray / bright yellow)."""
h, w = mask.shape
color = np.zeros((h, w, 3), dtype=np.uint8)
color[mask == 0] = COLOR_BACKGROUND
color[mask == 1] = COLOR_CROP
return Image.fromarray(color, mode="RGB")
def create_overlay(original: Image.Image, mask: np.ndarray) -> Image.Image:
"""Blend a semi-transparent magenta layer over predicted crop regions."""
base = np.asarray(original.convert("RGB"), dtype=np.float32)
crop_pixels = mask == 1
overlay = base.copy()
overlay[crop_pixels] = (
(1.0 - OVERLAY_ALPHA) * base[crop_pixels]
+ OVERLAY_ALPHA * np.array(OVERLAY_COLOR, dtype=np.float32)
)
return Image.fromarray(overlay.clip(0, 255).astype(np.uint8), mode="RGB")
def _load_font(size: int):
"""Load a default font, gracefully handling older Pillow versions."""
try:
return ImageFont.load_default(size=size)
except TypeError: # Pillow without the size argument
return ImageFont.load_default()
def combine_panels(
original: Image.Image,
color_mask: Image.Image,
overlay: Image.Image,
gap: int = 12,
label_height: int = 36,
) -> Image.Image:
"""
Combine the three panels horizontally into a single labeled image.
All panels are resized to identical dimensions (the original image size)
before being concatenated.
"""
target_size = original.size # (w, h)
panels = [
("Original", original.resize(target_size)),
("Predicted Mask", color_mask.resize(target_size)),
("Crop Overlay", overlay.resize(target_size)),
]
w, h = target_size
total_w = w * 3 + gap * 2
total_h = h + label_height
canvas = Image.new("RGB", (total_w, total_h), (255, 255, 255))
draw = ImageDraw.Draw(canvas)
font = _load_font(20)
x = 0
for label, img in panels:
canvas.paste(img, (x, label_height))
try:
text_w = draw.textlength(label, font=font)
except AttributeError: # very old Pillow
text_w = len(label) * 6
draw.text(
(x + (w - text_w) / 2, max(0, (label_height - 20) / 2)),
label,
fill=(0, 0, 0),
font=font,
)
x += w + gap
return canvas
def image_to_png_bytes(image: Image.Image) -> bytes:
"""Encode a PIL image as PNG bytes (used by the download button)."""
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return buffer.getvalue()
|