cottongin/pi-weekly-newspaper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement mosaic cover generation logic

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

131
src/cover.py
View File

@@ -1,52 +1,13 @@
import logging
import os
from collections import Counter
from datetime import date
from io import BytesIO
from urllib.parse import quote
import requests
from PIL import Image as PILImage, ImageDraw, ImageFont
import config
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:
try:
@@ -62,14 +23,6 @@ def _get_font(size: int) -> ImageFont.FreeTypeFont:
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(
text: str, font: ImageFont.FreeTypeFont, max_width: int
) -> str:
@@ -152,46 +105,64 @@ def generate_programmatic_cover(
return path
def generate_ai_cover(
def generate_mosaic_cover(
output_dir: str,
week_start: date,
week_end: date,
headlines: list[str],
categories: list[str],
image_paths: list[str],
) -> str:
os.makedirs(output_dir, exist_ok=True)
w, h = config.COVER_SIZE
img = PILImage.new("RGBA", (w, h), color=(245, 245, 245, 255))
dominant = _get_dominant_category(categories)
prompt = CATEGORY_PROMPTS.get(dominant, CATEGORY_PROMPTS["Default"])
url = config.POLLINATIONS_URL.format(prompt=quote(prompt))
if len(image_paths) == 1:
# Single image case
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:
response = requests.get(url, timeout=60)
response.raise_for_status()
for path in image_paths:
if col_heights[0] > h and col_heights[1] > h:
break # Both columns full
bg = PILImage.open(BytesIO(response.content))
if bg.mode not in ("RGB", "RGBA"):
bg = bg.convert("RGB")
bg = bg.resize((w, h), PILImage.Resampling.LANCZOS)
try:
with PILImage.open(path) as source_img:
source_img = source_img.convert("RGBA")
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")
overlay = PILImage.new("RGBA", (w, h), (0, 0, 0, 0))
draw = ImageDraw.Draw(overlay)
_draw_text_overlays(draw, w, h, week_start, week_end, headlines)
img = PILImage.alpha_composite(img, overlay)
col_idx = 0 if col_heights[0] <= col_heights[1] else 1
x = 4 + col_idx * (col_width + 4)
y = col_heights[col_idx]
out = img.convert("RGB")
filename = f"cover-{week_start.isoformat()}-ai.jpg"
path = os.path.join(output_dir, filename)
out.save(path, format="JPEG", progressive=False, quality=90)
return path
img.paste(resized, (x, y), resized)
col_heights[col_idx] += new_h + 4
except Exception as e:
logger.error("Failed to process image %s for mosaic: %s", path, e)
except Exception as e:
logger.error("AI cover generation failed, falling back to programmatic: %s", e)
return generate_programmatic_cover(
output_dir, week_start, week_end, headlines, categories
)
draw = ImageDraw.Draw(img)
_draw_text_overlays(draw, w, h, week_start, week_end, headlines)
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(
@@ -200,12 +171,14 @@ def generate_cover(
week_start: date,
week_end: date,
headlines: list[str],
categories: list[str],
image_paths: list[str],
) -> str:
if method == "ai":
return generate_ai_cover(
output_dir, week_start, week_end, headlines, categories
if method == "mosaic":
return generate_mosaic_cover(
output_dir, week_start, week_end, headlines, image_paths
)
# Fallback to programmatic text cover (ignore image_paths)
return generate_programmatic_cover(
output_dir, week_start, week_end, headlines, categories
output_dir, week_start, week_end, headlines, []
)

79
tests/test_cover.py
View File

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