synology-thumbgen/psthumbgen.py

import sys
import os
import re
import argparse
import errno
import time
import subprocess
from PIL import Image
from multiprocessing import Pool
from multiprocessing import Value
from PIL import Image, ImageDraw, ImageFont
import io
import struct


class State(object):
    def __init__(self):
        self.counter = Value('i', 0)
        self.start_ticks = Value('d', time.process_time())

    def increment(self, n=1):
        with self.counter.get_lock():
            self.counter.value += n

    @property
    def value(self):
        return self.counter.value

    @property
    def start(self):
        return self.start_ticks.value


def init(args):
    global state
    state = args


def main():
    args = parse_args()
    state = State()

    files = find_files(args.directory)

    with Pool(processes=4, initializer=init, initargs=(state, )) as pool:
        pool.map(process_file, files)

    print("{0} files processed in total.".format(state.value))


def parse_args():
    parser = argparse.ArgumentParser(
        description="Create thumbnails for Synology Photo Station.")
    parser.add_argument("--directory", required=True,
                        help="Directory to generate thumbnails for. "
                             "Subdirectories will always be processed.")

    return parser.parse_args()


def find_files(dir):
    # Only process formats that Synology doesn't handle well
    # Exclude common images (jpg, png, gif, etc.) since NAS handles them fine
    valid_exts = ('mp4', 'webm', 'avi', 'mkv', 'mov', 'wmv', 'flv', 'm4v', 'ts',
                  'psd', 'blend')
    valid_exts_re = "|".join(
        map((lambda ext: ".*\\.{0}$".format(ext)), valid_exts))

    for root, dirs, files in os.walk(dir):
        for name in files:
            if re.match(valid_exts_re, name, re.IGNORECASE) \
                    and not name.startswith('SYNOPHOTO_THUMB') \
                    and not name.startswith('SYNOVIDEO_VIDEO_SCREENSHOT'):
                yield os.path.join(root, name)


def print_progress():
    global state
    state.increment(1)
    processed = state.value
    if processed % 10 == 0:
        elapsed = float(time.process_time() - state.start)
        rate = float(processed) / elapsed if elapsed > 0 else float(processed)
        print("{0} files processed so far, averaging {1:.2f} files per second."
              .format(processed, rate))


def process_file(file_path):
    print(file_path)

    (dir, filename) = os.path.split(file_path)
    thumb_dir = os.path.join(dir, 'eaDir_tmp', filename)
    ensure_directory_exists(thumb_dir)

    # Determine file type and process accordingly
    file_ext = os.path.splitext(filename)[1].lower()
    
    if file_ext in ['.mp4', '.webm', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.m4v', '.ts']:
        create_video_thumbnails(file_path, thumb_dir)
    elif file_ext in ['.psd']:
        create_psd_thumbnails(file_path, thumb_dir)
    elif file_ext in ['.blend']:
        create_blend_thumbnails(file_path, thumb_dir)
    else:
        # Standard image processing
        create_thumbnails(file_path, thumb_dir)

    print_progress()


def ensure_directory_exists(path):
    try:
        os.makedirs(path)
    except OSError as exception:
        if exception.errno != errno.EEXIST:
            raise


def create_video_thumbnails(source_path, dest_dir):
    """Generate video thumbnails using FFmpeg"""
    to_generate = (('SYNOVIDEO_VIDEO_SCREENSHOT.jpg', 1280),
                   ('SYNOPHOTO_THUMB_XL.jpg', 1280),
                   ('SYNOPHOTO_THUMB_B.jpg', 640),
                   ('SYNOPHOTO_THUMB_M.jpg', 320),
                   ('SYNOPHOTO_THUMB_PREVIEW.jpg', 160),
                   ('SYNOPHOTO_THUMB_S.jpg', 120))

    for thumb in to_generate:
        thumb_path = os.path.join(dest_dir, thumb[0])
        if os.path.exists(thumb_path):
            continue

        try:
            # Use FFmpeg to extract a frame from the video at 10% duration
            cmd = [
                'ffmpeg', '-i', source_path,
                '-ss', '00:00:05',  # Seek to 5 seconds
                '-vframes', '1',
                '-vf', f'scale={thumb[1]}:{thumb[1]}:force_original_aspect_ratio=decrease',
                '-y',  # Overwrite output file
                thumb_path
            ]
            
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode != 0:
                print(f"Warning: FFmpeg failed for {source_path}: {result.stderr}")
                continue
                
        except FileNotFoundError:
            print("Warning: FFmpeg not found. Video thumbnails will be skipped.")
            print("Install FFmpeg: https://ffmpeg.org/download.html")
            break
        except Exception as e:
            print(f"Error processing video {source_path}: {e}")


def create_psd_thumbnails(source_path, dest_dir):
    """Generate PSD thumbnails using PIL with PSD support"""
    try:
        # Try to open PSD file - requires pillow-psd plugin
        from psd_tools import PSDImage
        psd = PSDImage.open(source_path)
        pil_image = psd.composite()
        
        to_generate = (('SYNOPHOTO_THUMB_XL.jpg', 1280),
                       ('SYNOPHOTO_THUMB_B.jpg', 640),
                       ('SYNOPHOTO_THUMB_M.jpg', 320),
                       ('SYNOPHOTO_THUMB_PREVIEW.jpg', 160),
                       ('SYNOPHOTO_THUMB_S.jpg', 120))

        for thumb in to_generate:
            thumb_path = os.path.join(dest_dir, thumb[0])
            if os.path.exists(thumb_path):
                continue

            pil_image.thumbnail((thumb[1], thumb[1]), Image.LANCZOS)
            # Convert to RGB if needed for JPEG
            if pil_image.mode == 'RGBA':
                pil_image = pil_image.convert('RGB')
            pil_image.save(thumb_path, 'JPEG')
            
    except ImportError:
        print("Warning: psd-tools not installed. Install with: pip install psd-tools")
    except Exception as e:
        print(f"Error processing PSD {source_path}: {e}")


def extract_blend_thumbnail(blend_path):
    """Extracts the largest PNG thumbnail from all PREV blocks in a .blend file."""
    import struct, re
    with open(blend_path, 'rb') as f:
        data = f.read()

    if not data.startswith(b'BLENDER'):
        return None
    pointer_size = 8 if data[7:8] == b'-' else 4
    endian = '<' if data[8:9] == b'v' else '>'
    offset = 12
    candidates = []
    while offset + 16 < len(data):
        code = data[offset:offset+4]
        size = struct.unpack(endian + 'I', data[offset+4:offset+8])[0]
        block_data_start = offset + 16 + pointer_size
        block_data_end = block_data_start + size
        if code == b'PREV':
            block = data[block_data_start:block_data_end]
            # Find all PNGs in this block
            for m in re.finditer(b'\x89PNG\r\n\x1a\n', block):
                start = m.start()
                end = block.find(b'IEND', start)
                if end != -1:
                    end += 8
                    png_data = block[start:end]
                    try:
                        img = Image.open(io.BytesIO(png_data))
                        candidates.append((img.size[0] * img.size[1], img))
                    except Exception:
                        continue
        offset = block_data_end
    if not candidates:
        return None
    return max(candidates, key=lambda x: x[0])[1]


def create_blend_thumbnails(source_path, dest_dir):
    """Extract embedded Blender thumbnail and generate Synology thumbnails."""
    img = extract_blend_thumbnail(source_path)
    if img is None:
        print(f"No embedded thumbnail in {source_path}")
        return

    to_generate = (('SYNOPHOTO_THUMB_XL.jpg', 1280),
                   ('SYNOPHOTO_THUMB_B.jpg', 640),
                   ('SYNOPHOTO_THUMB_M.jpg', 320),
                   ('SYNOPHOTO_THUMB_PREVIEW.jpg', 160),
                   ('SYNOPHOTO_THUMB_S.jpg', 120))

    for thumb in to_generate:
        thumb_path = os.path.join(dest_dir, thumb[0])
        if os.path.exists(thumb_path):
            continue
        im_copy = img.copy()
        im_copy.thumbnail((thumb[1], thumb[1]), Image.LANCZOS)
        if im_copy.mode == 'RGBA':
            im_copy = im_copy.convert('RGB')
        im_copy.save(thumb_path, 'JPEG', quality=85)


def create_thumbnails(source_path, dest_dir):
    """Original image thumbnail creation"""
    try:
        im = Image.open(source_path)

        to_generate = (('SYNOPHOTO_THUMB_XL.jpg', 1280),
                       ('SYNOPHOTO_THUMB_B.jpg', 640),
                       ('SYNOPHOTO_THUMB_M.jpg', 320),
                       ('SYNOPHOTO_THUMB_PREVIEW.jpg', 160),
                       ('SYNOPHOTO_THUMB_S.jpg', 120))

        for thumb in to_generate:
            thumb_path = os.path.join(dest_dir, thumb[0])
            if os.path.exists(thumb_path):
                continue

            # Create a copy for thumbnailing
            im_copy = im.copy()
            im_copy.thumbnail((thumb[1], thumb[1]), Image.LANCZOS)
            im_copy.save(thumb_path)
            
    except Exception as e:
        print(f"Error processing image {source_path}: {e}")


if __name__ == "__main__":
    sys.exit(main())