feat: implement mosaic cover generation logic

Made-with: Cursor
This commit is contained in:
cottongin
2026-04-06 19:21:55 -04:00
parent b2d084cd69
commit 362e0c9b8e
2 changed files with 73 additions and 137 deletions

View File

@@ -1,52 +1,13 @@
import logging import logging
import os import os
from collections import Counter
from datetime import date from datetime import date
from io import BytesIO
from urllib.parse import quote
import requests
from PIL import Image as PILImage, ImageDraw, ImageFont from PIL import Image as PILImage, ImageDraw, ImageFont
import config import config
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
CATEGORY_PROMPTS = {
"Government": (
"Watercolor painting of Plymouth Massachusetts town hall, "
"New England civic architecture, soft morning light, no text"
),
"Culture": (
"Plymouth Massachusetts waterfront boardwalk, autumn festival "
"atmosphere, warm tones, New England charm, no text"
),
"Culture Calendar": (
"Plymouth Massachusetts harbor at golden hour, sailboats and "
"historic waterfront, warm autumn palette, no text"
),
"Sports": (
"Plymouth Massachusetts harbor marina at sunrise, boats and "
"morning mist on Cape Cod Bay, no text"
),
"Environment": (
"Cape Cod Bay shoreline near Plymouth Massachusetts, coastal "
"wetlands and dune grass, golden hour light, no text"
),
"Opinion": (
"Plymouth Rock historic site in Plymouth Massachusetts, "
"dramatic sky, iconic New England landmark, no text"
),
"Feature": (
"Downtown Plymouth Massachusetts street scene, historic "
"brick storefronts, New England small town charm, no text"
),
"Default": (
"Aerial view of Plymouth Massachusetts coastal town, serene "
"harbor, Cape Cod Bay, New England charm, no text"
),
}
def _get_font(size: int) -> ImageFont.FreeTypeFont: def _get_font(size: int) -> ImageFont.FreeTypeFont:
try: try:
@@ -62,14 +23,6 @@ def _get_font(size: int) -> ImageFont.FreeTypeFont:
return ImageFont.load_default() return ImageFont.load_default()
def _get_dominant_category(categories: list[str]) -> str:
if not categories:
return "Default"
counter = Counter(categories)
most_common = counter.most_common(1)[0][0]
return most_common
def _truncate_to_width( def _truncate_to_width(
text: str, font: ImageFont.FreeTypeFont, max_width: int text: str, font: ImageFont.FreeTypeFont, max_width: int
) -> str: ) -> str:
@@ -152,46 +105,64 @@ def generate_programmatic_cover(
return path return path
def generate_ai_cover( def generate_mosaic_cover(
output_dir: str, output_dir: str,
week_start: date, week_start: date,
week_end: date, week_end: date,
headlines: list[str], headlines: list[str],
categories: list[str], image_paths: list[str],
) -> str: ) -> str:
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
w, h = config.COVER_SIZE w, h = config.COVER_SIZE
img = PILImage.new("RGBA", (w, h), color=(245, 245, 245, 255))
dominant = _get_dominant_category(categories) if len(image_paths) == 1:
prompt = CATEGORY_PROMPTS.get(dominant, CATEGORY_PROMPTS["Default"]) # Single image case
url = config.POLLINATIONS_URL.format(prompt=quote(prompt)) try:
with PILImage.open(image_paths[0]) as source_img:
source_img = source_img.convert("RGBA")
# Resize to fit within 480x800 preserving aspect ratio
ratio = min(w / source_img.width, h / source_img.height)
new_size = (int(source_img.width * ratio), int(source_img.height * ratio))
resized = source_img.resize(new_size, PILImage.Resampling.LANCZOS)
# Center it
x = (w - new_size[0]) // 2
y = (h - new_size[1]) // 2
img.paste(resized, (x, y), resized)
except Exception as e:
logger.error("Failed to process single image for cover: %s", e)
elif len(image_paths) > 1:
# Mosaic case
col_width = (w - (3 * 4)) // 2
col_heights = [4, 4]
try: for path in image_paths:
response = requests.get(url, timeout=60) if col_heights[0] > h and col_heights[1] > h:
response.raise_for_status() break # Both columns full
bg = PILImage.open(BytesIO(response.content)) try:
if bg.mode not in ("RGB", "RGBA"): with PILImage.open(path) as source_img:
bg = bg.convert("RGB") source_img = source_img.convert("RGBA")
bg = bg.resize((w, h), PILImage.Resampling.LANCZOS) new_h = int(source_img.height * (col_width / source_img.width))
resized = source_img.resize((col_width, new_h), PILImage.Resampling.LANCZOS)
img = bg.convert("RGBA") col_idx = 0 if col_heights[0] <= col_heights[1] else 1
overlay = PILImage.new("RGBA", (w, h), (0, 0, 0, 0)) x = 4 + col_idx * (col_width + 4)
draw = ImageDraw.Draw(overlay) y = col_heights[col_idx]
_draw_text_overlays(draw, w, h, week_start, week_end, headlines)
img = PILImage.alpha_composite(img, overlay)
out = img.convert("RGB") img.paste(resized, (x, y), resized)
filename = f"cover-{week_start.isoformat()}-ai.jpg" col_heights[col_idx] += new_h + 4
path = os.path.join(output_dir, filename) except Exception as e:
out.save(path, format="JPEG", progressive=False, quality=90) logger.error("Failed to process image %s for mosaic: %s", path, e)
return path
except Exception as e: draw = ImageDraw.Draw(img)
logger.error("AI cover generation failed, falling back to programmatic: %s", e) _draw_text_overlays(draw, w, h, week_start, week_end, headlines)
return generate_programmatic_cover(
output_dir, week_start, week_end, headlines, categories out = img.convert("RGB")
) filename = f"cover-{week_start.isoformat()}-mosaic.jpg"
path = os.path.join(output_dir, filename)
out.save(path, format="JPEG", progressive=False, quality=90)
return path
def generate_cover( def generate_cover(
@@ -200,12 +171,14 @@ def generate_cover(
week_start: date, week_start: date,
week_end: date, week_end: date,
headlines: list[str], headlines: list[str],
categories: list[str], image_paths: list[str],
) -> str: ) -> str:
if method == "ai": if method == "mosaic":
return generate_ai_cover( return generate_mosaic_cover(
output_dir, week_start, week_end, headlines, categories output_dir, week_start, week_end, headlines, image_paths
) )
# Fallback to programmatic text cover (ignore image_paths)
return generate_programmatic_cover( return generate_programmatic_cover(
output_dir, week_start, week_end, headlines, categories output_dir, week_start, week_end, headlines, []
) )

View File

@@ -1,19 +1,16 @@
import json import json
import os import os
from collections import Counter
from datetime import date from datetime import date
from unittest.mock import patch, MagicMock from unittest.mock import patch
from io import BytesIO from io import BytesIO
from PIL import Image as PILImage, ImageDraw, ImageFont from PIL import Image as PILImage, ImageDraw, ImageFont
from src.cover import ( from src.cover import (
generate_programmatic_cover, generate_programmatic_cover,
generate_ai_cover, generate_mosaic_cover,
generate_cover, generate_cover,
_get_dominant_category,
_truncate_to_width, _truncate_to_width,
_get_font, _get_font,
CATEGORY_PROMPTS,
) )
@@ -46,54 +43,31 @@ def test_programmatic_cover_no_headlines(tmp_path):
assert img.size == (480, 800) assert img.size == (480, 800)
def test_ai_cover_is_portrait(tmp_path): def test_generate_mosaic_cover_single_image(tmp_path):
fake_img = PILImage.new("RGB", (480, 800), color="blue") # Create a dummy image
buf = BytesIO() img_path = str(tmp_path / "dummy.jpg")
fake_img.save(buf, format="JPEG") img = PILImage.new("RGB", (1000, 500), color="red")
fake_bytes = buf.getvalue() img.save(img_path)
mock_response = MagicMock() output_dir = str(tmp_path)
mock_response.content = fake_bytes week_start = date(2026, 4, 6)
mock_response.raise_for_status = MagicMock() week_end = date(2026, 4, 12)
headlines = ["Test Headline"]
with patch("src.cover.requests.get", return_value=mock_response): cover_path = generate_mosaic_cover(output_dir, week_start, week_end, headlines, [img_path])
path = generate_ai_cover(
output_dir=str(tmp_path),
week_start=date(2026, 4, 6),
week_end=date(2026, 4, 12),
headlines=["Test Headline"],
categories=["Government"],
)
assert os.path.exists(path) assert os.path.exists(cover_path)
img = PILImage.open(path) with PILImage.open(cover_path) as result_img:
assert img.format == "JPEG" assert result_img.size == (480, 800)
assert img.size == (480, 800)
def test_ai_cover_falls_back_to_programmatic(tmp_path):
with patch("src.cover.requests.get", side_effect=Exception("API down")):
path = generate_ai_cover(
output_dir=str(tmp_path),
week_start=date(2026, 4, 6),
week_end=date(2026, 4, 12),
headlines=["Test"],
categories=["Government"],
)
assert os.path.exists(path)
img = PILImage.open(path)
assert img.format == "JPEG"
assert img.size == (480, 800)
def test_generate_cover_dispatches(tmp_path): def test_generate_cover_dispatches(tmp_path):
with patch("src.cover.generate_ai_cover") as mock_ai: with patch("src.cover.generate_mosaic_cover") as mock_mosaic:
mock_ai.return_value = "/fake/ai.jpg" mock_mosaic.return_value = "/fake/mosaic.jpg"
result = generate_cover("ai", str(tmp_path), date(2026, 4, 6), result = generate_cover("mosaic", str(tmp_path), date(2026, 4, 6),
date(2026, 4, 12), ["A"], ["Government"]) date(2026, 4, 12), ["A"], ["/x/a.jpg"])
assert result == "/fake/ai.jpg" assert result == "/fake/mosaic.jpg"
mock_ai.assert_called_once() mock_mosaic.assert_called_once()
with patch("src.cover.generate_programmatic_cover") as mock_prog: with patch("src.cover.generate_programmatic_cover") as mock_prog:
mock_prog.return_value = "/fake/text.jpg" mock_prog.return_value = "/fake/text.jpg"
@@ -103,17 +77,6 @@ def test_generate_cover_dispatches(tmp_path):
mock_prog.assert_called_once() mock_prog.assert_called_once()
def test_get_dominant_category():
assert _get_dominant_category(["Government", "Government", "Culture"]) == "Government"
assert _get_dominant_category(["Culture", "Government"]) == "Culture"
assert _get_dominant_category([]) == "Default"
def test_category_prompts_include_massachusetts():
for category, prompt in CATEGORY_PROMPTS.items():
assert "Massachusetts" in prompt, f"Prompt for {category} missing 'Massachusetts'"
def test_truncate_to_width_respects_pixel_budget(): def test_truncate_to_width_respects_pixel_budget():
font = _get_font(18) font = _get_font(18)
short = "Short text" short = "Short text"