Coverage for imcluster/thumbnails.py: 100.00%

1"""Thumbnail generation for cluster reports.""" 

2 

3import base64 

4from io import BytesIO 

5from pathlib import Path 

6 

7from PIL import Image, ImageOps, UnidentifiedImageError 

8from rich.console import Console 

9 

10from .io import ImclusterIO 

11 

12console = Console() 

13 

14 

15def generate_thumbnail(path: str | Path, width: int, height: int) -> str: 

16 """Create a base64-encoded JPEG thumbnail bounded by given dimensions. 

17 

18 Args: 

19 path: Source image path. 

20 width: Maximum thumbnail width in pixels. 

21 height: Maximum thumbnail height in pixels. 

22 

23 Returns: 

24 ASCII base64 data for the generated JPEG. 

25 """ 

26 try: 

27 with Image.open(path) as source: 

28 image = ImageOps.exif_transpose(source).convert("RGB") 

29 image.thumbnail((width, height), Image.Resampling.LANCZOS) 

30 buffered = BytesIO() 

31 image.save(buffered, format="JPEG") 

32 except (OSError, UnidentifiedImageError) as error: 

33 raise ValueError(f"Cannot create thumbnail for '{path}': {error}") from error 

34 return base64.b64encode(buffered.getvalue()).decode("ascii") 

35 

36 

37def generate_thumbnails( 

38 imcluster_io: ImclusterIO, 

39 thumbnail_width: int = 256, 

40 thumbnail_height: int = 256, 

41 force: bool = False, 

42 force_thumbnails: bool = False, 

43) -> None: 

44 """Generate and cache thumbnails used by the HTML cluster report. 

45 

46 Args: 

47 imcluster_io: Image collection and its persisted result table. 

48 thumbnail_width: Maximum thumbnail width in pixels. 

49 thumbnail_height: Maximum thumbnail height in pixels. 

50 force: Regenerate thumbnails regardless of cached data. 

51 force_thumbnails: Regenerate only the thumbnail cache. 

52 """ 

53 

54 if not imcluster_io.has_column("thumbnail") or force or force_thumbnails: 

55 console.print( 

56 "[cyan]Generating thumbnails:[/cyan] " 

57 f"{len(imcluster_io.images)} images, maximum size " 

58 f"{thumbnail_width}x{thumbnail_height}; caching results in " 

59 f"'{imcluster_io.output}'." 

60 ) 

61 with console.status("[cyan]Creating and caching thumbnails...[/cyan]"): 

62 imcluster_io.save_column( 

63 "thumbnail", 

64 imcluster_io.df.apply( 

65 lambda row: generate_thumbnail( 

66 row["path"], 

67 thumbnail_width, 

68 thumbnail_height, 

69 ), 

70 axis=1, 

71 ), 

72 ) 

73 else: 

74 console.print( 

75 "[green]Using cached thumbnails:[/green] " 

76 f"loaded {len(imcluster_io.images)} thumbnails from " 

77 f"'{imcluster_io.output}'." 

78 )