add video and audio thumbnail and preview generation, add 360 degree support, improve mesh image quality

Author: Windowsfreak

image_processor.py

@@ -1,82 +1,231 @@
+import base64
 import io
 import os
 import sys
 import json
 import pyvips
 import numpy as np
-
+import subprocess
+import tempfile
+import xml.etree.ElementTree as et
 
 def resize_image(image, width, height):
     # Use the smaller scaling factor to maintain aspect ratio
     scale_x = width / image.width
     scale_y = height / image.height
     scale = min(scale_x, scale_y)
-
     # Resize the image using libvips if needed
     if scale < 1:
         return image.resize(scale)
     else:
         return image
 
-
 def resize_fixed(image, width, height):
     # Use the smaller scaling factor to maintain aspect ratio
     scale_x = width / image.width
     scale_y = height / image.height
-
     # Resize the image using libvips
     return image.resize(scale_x, vscale=scale_y, gap=8)
 
-
 def mesh_gradient(image):
-    glyphs = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
-    text = ""
+    red_levels = np.array([0, 51, 102, 153, 204, 255])
+    green_levels = np.array([0, 36, 73, 109, 146, 182, 219, 255])
+    blue_levels = np.array([0, 64, 128, 192, 255])
+
+    palette = np.array([(r, g, b) for r in red_levels for g in green_levels for b in blue_levels], dtype=np.uint8)
 
     array = image.write_to_memory()
     np_array = np.frombuffer(array, dtype=np.uint8)
-
     np_array = np_array.reshape(image.height, image.width, image.bands)
 
-    for y in range(4):
-        for x in range(4):
-            r, g, b = np_array[y, x, :3]
-            color_val = ((r & 0xF0) << 4) | (g & 0xF0) | ((b & 0xF0) >> 4)
-            text += glyphs[color_val >> 6]
-            text += glyphs[color_val & 63]
-
-    return text
-
+    indices = []
+    for y in range(8):
+        for x in range(8):
+            pixel = np_array[y, x, :3]
+            # Calculate squared Euclidean distance to each palette color
+            diffs = palette.astype(np.int32) - pixel.astype(np.int32)
+            dists = np.sum(diffs ** 2, axis=1)
+            idx = np.argmin(dists)
+            indices.append(idx)
+
+    byte_seq = bytes(indices)
+    return base64.b64encode(byte_seq).decode('ascii')
+
+def get_duration(input_file):
+    cmd = [
+        'ffprobe', '-v', 'error', '-show_entries', 'format=duration',
+        '-of', 'default=noprint_wrappers=1:nokey=1', input_file
+    ]
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        return 0
+    try:
+        return int(float(result.stdout.strip()))
+    except ValueError:
+        return 0
+
+def generate_video_frame(input_file, duration, temp_dir):
+    seek_time = duration * 0.1
+    output_jxl = os.path.join(temp_dir, 'frame.jxl')
+    cmd = [
+        'ffmpeg', '-i', input_file, '-ss', str(seek_time), '-vframes', '1',
+        '-q:v', '100', output_jxl
+    ]
+    subprocess.run(cmd, check=True)
+    return output_jxl
+
+def generate_audio_waveform(input_file, temp_dir):
+    output_png = os.path.join(temp_dir, 'waveform.png')
+    cmd = [
+        'ffmpeg', '-i', input_file, '-f', 'lavfi', '-i', 'color=c=#000000:s=640x120',
+        '-filter_complex',
+        '[0:a] aformat=channel_layouts=mono,showwavespic=s=640x120:colors=#808080:filter=peak:scale=sqrt [pk]; '
+        '[0:a] aformat=channel_layouts=mono,showwavespic=s=640x120:colors=#ffffff:scale=sqrt [rms], '
+        '[pk] [rms] overlay=format=auto [nobg], [1:v] [nobg] overlay=format=auto',
+        '-frames:v', '1', '-update', 'true', output_png
+    ]
+    subprocess.run(cmd, check=True)
+    return output_png
+
+def generate_video_preview(input_file, output_file):
+    # Step 1: Get input framerate using ffprobe
+    ffprobe_cmd = [
+        'ffprobe', '-v', 'error', '-select_streams', 'v:0',
+        '-show_entries', 'stream=avg_frame_rate', '-of',
+        'default=noprint_wrappers=1:nokey=1', input_file
+    ]
+    result = subprocess.run(ffprobe_cmd, capture_output=True, text=True, check=True)
+    framerate_str = result.stdout.strip()  # e.g., "30/1"
+    if not framerate_str:
+        framerate_str = "1/1"
+
+    # Parse framerate to float (e.g., "30/1" -> 30.0)
+    num, den = map(int, framerate_str.split('/'))
+    input_fps = num / den
+
+    # Step 2: Build FFmpeg command
+    preview_file = output_file + '.mkv'
+    cmd = [
+        'ffmpeg', '-i', input_file, '-c:v', 'libsvtav1', '-preset', '5',
+        '-crf', '56', '-profile:v', 'main', '-level:v', '5.1',
+        '-c:a', 'libopus', '-b:a', '16k', '-ac', '1', '-vbr', 'on'
+    ]
+    vf_filters = ['scale=480:-1']
+    if input_fps > 30:
+        vf_filters.insert(0, 'fps=fps=source_fps/2')
+    cmd.extend(['-vf', ','.join(vf_filters)])
+    cmd.append(preview_file)
+
+    # Run the command
+    subprocess.run(cmd, check=True)
+
+def generate_audio_preview(input_file, output_file):
+    preview_file = output_file + '.mkv'
+    cmd = [
+        'ffmpeg', '-i', input_file, '-c:a', 'libopus', '-b:a', '16k',
+        '-ac', '1', '-vbr', 'on', '-vn', preview_file
+    ]
+    subprocess.run(cmd, check=True)
+
+def is_spherical(image):
+    if image.width / image.height != 2:
+        return False
+    xmp_bytes = None
+    try:
+        if 'xmp-data' in image.get_fields():
+            xmp_bytes = image.get('xmp-data')
+        elif 'xmp' in image.get_fields():
+            xmp_bytes = image.get('xmp')
+    except Exception:
+        return False
+    if not xmp_bytes:
+        return False
+    namespaces = {'x': 'adobe:ns:meta/', 'rdf': 'http://www.w3.org/1999/02/22-rdf-syntax-ns#', 'GPano': 'http://ns.google.com/photos/1.0/panorama/'}
+    try:
+        xmp_str = xmp_bytes.decode('utf-8')
+        root = et.fromstring(xmp_str)
+        rdf_descr = root.find('.//rdf:Description', namespaces)
+        if rdf_descr is not None:
+            proj = rdf_descr.attrib.get('{http://ns.google.com/photos/1.0/panorama/}ProjectionType')
+            if proj and proj == 'equirectangular':
+                return True
+    except Exception:
+        pass
+    return False
+
+def generate_cubemap(input_file, output_file, size):
+    cubemap_file = output_file + '.c.jxl'
+    cmd = ['kubi', '-l', 'row', '-s', str(size), '--order', '1', '4', '0', '5', '3', '2', input_file, cubemap_file]  # Adjust if kubi args differ
+    subprocess.run(cmd, check=True)
 
 def process_image(image_info):
-    if image_info['input_file'].lower().endswith((".heif", ".heic")):
-        image = pyvips.Image.new_from_file(image_info['input_file'], memory=True, unlimited=True)
+    input_file = image_info['input_file']
+    output_file = image_info['output_file']
+    orientation = image_info.get('orientation', 0)
+
+    _, ext = os.path.splitext(input_file.lower())
+
+    video_exts = [".3gp", ".flv", ".mov", ".qt", ".m2ts", ".mts", ".divx", ".vob", ".webm", ".mkv", ".mka", ".wmv", ".avi", ".mp4", ".mpg", ".mpeg", ".ps", ".ts", ".rm", ".ogv", ".dv"]
+    audio_exts = [".mp3", ".wav", ".opus", ".aac", ".ogg", ".wma", ".m4a", ".flac", ".alac", ".mka"]
+    is_video = ext in video_exts
+    is_audio = ext in audio_exts and not is_video
+
+    duration = None
+    if is_video or is_audio:
+        duration = get_duration(input_file)
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            if is_video:
+                temp_file = generate_video_frame(input_file, duration, temp_dir)
+                generate_video_preview(input_file, output_file)
+            else:
+                temp_file = generate_audio_waveform(input_file, temp_dir)
+                generate_audio_preview(input_file, output_file)
+
+            # Load the generated thumbnail for processing
+            image = pyvips.Image.new_from_file(temp_file, memory=True)
     else:
-        image = pyvips.Image.new_from_file(image_info['input_file'], memory=True)
+        if ext in (".heif", ".heic"):
+            image = pyvips.Image.new_from_file(input_file, memory=True, unlimited=True)
+        else:
+            image = pyvips.Image.new_from_file(input_file, memory=True)
 
-    if image_info['orientation'] > 1:
+    if orientation > 1:
         image = image.autorot()
 
     true_width = image.width
     true_height = image.height
 
-    image.jxlsave(image_info['output_file'] + '.o.jxl', Q=75, strip=True, effort=4)
+    is_sph = False if is_video or is_audio else is_spherical(image)
+    if is_sph:
+        generate_cubemap(input_file, output_file, min(true_width / 4, 1024))
+
+    image.jxlsave(output_file + '.o.jxl', Q=75, strip=True, effort=4)
+
     image = resize_image(image, 2048, 2048)
-    image.jxlsave(image_info['output_file'] + '.h.jxl', Q=60, strip=True, effort=5)
+    image.jxlsave(output_file + '.h.jxl', Q=60, strip=True, effort=5)
+
     image = resize_image(image, 400, 200)
-    image.jxlsave(image_info['output_file'] + '.s.jxl', Q=20, strip=True, effort=5)
-    image = resize_fixed(image, 4, 4)
+    image.jxlsave(output_file + '.s.jxl', Q=20, strip=True, effort=5)
+
+    image = resize_fixed(image, 8, 8)
     image = image.colourspace("srgb")
 
-    return {
+    result = {
         "width": true_width,
         "height": true_height,
         "color": mesh_gradient(image)
     }
+    if is_sph:
+        result["pano"] = True
+    if duration is not None:
+        result["duration"] = duration
 
+    return result
 
 if __name__ == "__main__":
     pyvips.voperation.cache_set_max_mem(2048)
     stdin_utf8 = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8')
     stdin_data = stdin_utf8.read()
     data = json.loads(stdin_data)
-    sys.stdout.write(json.dumps(process_image(data)))
+    sys.stdout.write(json.dumps(process_image(data)))