unscore - Script for converting score videos into playable PDFs. git clone https://benconnors.ca/git-repos/unscore

About | Log | Files | Refs

unscore.py (14999B) - raw

---

 #!/usr/bin/env python3
 
 import os
 import re
 import subprocess as subp
 import tempfile
 from typing import List
 
 import numpy as np
 from PIL import Image, ImageOps, ImageFilter, ImageMorph, ImageChops, ImageDraw, ImageFont
 try:
     import scipy.ndimage
     import scipy.signal
 
     HAVE_SCIPY = True
 except ModuleNotFoundError:
     HAVE_SCIPY = False
 
 ## BEGIN Thresholding
 
 def threshold(image: Image, level: int) -> Image:
     """Threshold an image at the given level."""
     return image.convert("L").point(lambda p: p >= level and 255, mode="1")
 
 def compute_otsu_level(image: Image) -> int:
     """Compute the Otsu level for thresholding."""
     hist = image.histogram()
     total = sum(hist)
 
     sum_b = 0
     w_b = 0
     mx = 0
     level = 0
     sum_1 = sum((a*b for a, b in zip(range(256), hist)))
     for i in range(1, 256):
         w_f = total - w_b
         if w_b > 0 and w_f > 0:
             m_f = (sum_1 - sum_b) / w_f
             val = w_b * w_f * ((sum_b / w_b) - m_f) * ((sum_b / w_b) - m_f)
             if val >= mx:
                 level = i
                 mx = val
         w_b += hist[i]
         sum_b += (i-1) * hist[i]
 
     return level
 
 def otsu(image: Image) -> Image:
     """Threshold an image using the Otsu level.
 
     Otsu thresholding typically does very well when the image is first upsampled, then    thresholded, e.g. via
 
       otsu(img.resize((img.width*2, img.height*2), resample=Image.Resampling.LANCZOS))
 
     This is useful when e.g. the original image is low-resolution; upsampling then
     thresholding will typically give a better result than thresholding and then
     upsampling.
     """
     image = image.convert("L")
     o = compute_otsu_level(image)
     return image.point(lambda i: i > o, mode="1")
 
 def local_otsu(image, window_size=100):
     """Threshold the image via the local Otsu method."""
     output = Image.new("1", image.size)
     for x_window in range(math.ceil(image.width / window_size)):
         for y_window in range(math.ceil(image.height / window_size)):
             a = x_window*window_size
             b = min(x_window*window_size + window_size, image.width)
             c = y_window*window_size
             d = min(y_window*window_size + window_size, image.height)
 
             output.paste(otsu(image.crop((a, c, b, d))), (a, c))
 
     return output
 
 ## END Thresholding
 
 ## BEGIN Image processing
 
 def difference(a: Image, b: Image, level: int = 50):
     """Compute a difference mask of two images.
 
     We do:
 
     1. Compute the difference between `a` and `b`
     2. Threshold this to remove noise using `level`
     3. Use grayscale dilation to expand this slightly
     4. Threshold
     """
     ## TODO: Is step 2 necessary? At worst we should just change some artifacting
     return ImageMorph.MorphOp(op_name="dilation8").apply(threshold(ImageChops.difference(a, b), level).convert("L"))[1].convert("1")
 
 def remove_line(a: Image, b: Image, c: Image, level: int = 50) -> Image:
     """Remove a moving line from a sequence of frames.
 
     The three images `a`, `b`, and `c` should each have the line in a different spot.
     We do:
 
     1. Compute the difference mask between `a` and `b`
     2. Paste these parts of image `c` into `a`.
 
     If the images are of varying quality, `a` should be the best.
     """
     a = a.copy()
     a.paste(c, None, mask=difference(a, b, level=level))
     
     return a
 
 def do_savgol(scene: Image, cutoff: float) -> Image:
     """Run a 2D Savitzky-Golay filter on the image, then threshold at `cutoff`.
 
     This is very good at smoothing out thresholded images and should be used as the
     final step; it is computationally expensive.
     """
     if not HAVE_SCIPY:
         raise ModuleNotFoundError("Need scipy for savgol filtering")
     a = np.array(scene).astype(float)
     b = scipy.signal.savgol_filter(
         scipy.signal.savgol_filter(a, 11, 4, axis=1), 
         11, 
         4, 
         axis=0
     )
     return Image.fromarray(b > cutoff)
 
 def split_lines(image: Image, min_dist: int = 15) -> List[Image]:
     """Split an image into lines, each separated by at least `min_dist`."""
     ## TODO: Morphological opening here to remove noise?
     image = otsu(image)
     arr = np.array(image)
 
     extents = []
     last = None
     for n, row in enumerate(arr):
         if False not in row:
             if last is None:
                 continue
             else:
                 extents.append((last, n))
                 last = None
         else:
             if last is None:
                 last = n
 
     if last is not None:
         extents.append((last, len(arr)))
 
     found = True
     while found:
         found = False
         for i in range(len(extents)-1):
             a, b = extents[i]
             a2, b2 = extents[i+1]
             if a2 - b < min_dist:
                 extents.pop(i)
                 extents.pop(i)
                 extents.insert(i, (a, b2))
                 found = True
                 break
 
     return extents
 
 ## END Image processing
 
 ## BEGIN Crop/pad
 
 def double_bbox(image: Image) -> [int, int, int, int]:
     """Get the bounding box of an image possibly surrounded by a black then white 
     outline.
     """
     image = threshold(image, 95)
     left, upper, right, lower = image.getbbox()
 
     other = image.crop((left, upper, right, lower))
     left2, upper2, right2, lower2 = ImageOps.invert(other).getbbox()
 
     return left2+left, upper2+upper, right2+left, lower2+upper
 
 def crop_scene(i: Image, new_width: int) -> Image:
     """Remove the bounding box from a scene and resize to match `new_width`."""
     width, height = i.width, i.height
     left, upper, right, lower = double_bbox(i)
     i = i.crop((left, upper, right, lower))
     
     new_height = int((lower - upper)*new_width/(right-left)+1)
     ri = i.resize((new_width, new_height))
 
     return ri
 
 def pad_scene(scene: Image, padding: int = 10) -> Image:
     """Add padding to a scene."""
     i = Image.new("L", size=(scene.width, scene.height+2*padding), color=255)
     i.paste(scene, (0, padding))
 
     return i
 
 ## END Crop/pad
 
 ## BEGIN FFmpeg
 
 def get_scenes(file: str, cutoff: float = 0.05) -> List[float]:
     """Use FFmpeg to get a list of scenes in a video file."""
     res = subp.check_output(["ffmpeg", "-i", file, "-filter:v", "select='gt(scene,%s)',showinfo" % str(cutoff), "-f", "null", "-"], stderr=subp.STDOUT)
 
     ## FFmpeg's output is annoying to parse...
     times = []
     for line in res.decode("utf-8").split('\n'):
         if not "showinfo" in line:
             continue
         val = re.search(r'pts_time:([0-9.]+)', line)
         if not val:
             continue
         times.append(float(val.group(1)))
 
     return [0] + times
 
 def do_screenshots(file: str, td: str, fps: int = 1) -> List[str]:
     """Use FFmpeg to take `fps` screenshots per second of `file`, storing them in 
     the folder `td`.
 
     Space-wise this is wasteful but is faster/simpler than taking only the screenshots
     we need.
     """
     subp.check_call(["ffmpeg", "-i", file, "-q:v", "3", "-vf", "fps=%d" % fps, os.path.join(td, "%09d.jpg")], stderr=subp.DEVNULL)
     snapshots = []
     for f in os.listdir(td):
         if not f.endswith(".jpg"):
             continue
         time = int(f.split('.')[0])
         snapshots.append((time, os.path.join(td, f)))
 
     return sorted(snapshots)
 
 ## END FFmpeg
 
 ## BEGIN Scenes
 
 def filter_scenes(scenes: List[float], cutoff: float = 1) -> List[float]:
     """Filter out scenes that are too close together (false positives)."""
     new = []
     last = None
     for s in scenes[::-1]:
         if last is None or last - s > cutoff:
             new.append(s)
             last = s
 
     return list(reversed(new))
 
 def find_scenes(f: str, ffmpeg_cutoff: float, line: bool = False, scene_cutoff: float = 1, fps: int = 1) -> List[Image]:
     """Return a list of scenes in the file `f`."""
     ## Find the times we need
     times = filter_scenes(get_scenes(f, ffmpeg_cutoff), scene_cutoff)
     with tempfile.TemporaryDirectory() as td:
         ## Take screenshots
         snapshots = do_screenshots(f, td, fps=fps)
 
         ## Open the screenshots we need
         use = []
         for n, (t, nt) in enumerate(zip(times, times[1:]+[None]), 1):
             ## Figure out which screenshots work for this scene
             available = [b for a, b in snapshots if t < a and (nt is None or a < nt)]
 
             if line:
                 ## Grab three shots to remove the line
                 p1 = Image.open(available[len(available)//3])
                 p2 = Image.open(available[len(available)//2])
                 p3 = Image.open(available[2*(len(available)//3)])
 
                 use.append(remove_line(p2, p1, p3))
             else:
                 ## Grab the middle, usually the compression is better
                 pick =  available[len(available)//2]
                 use.append(Image.open(pick))
 
         return use
 
 def split_scene(scene: Image, min_dist: int = 10) -> List[Image]:
     """Split a scene into lines."""
     todo = []
     extents = split_lines(scene, min_dist)
     for a, b in extents:
         todo.append(scene.crop((0, a, scene.width, b)))
     return todo
 
 def merge_lines(images: List[Image], aspect: float = 8.5/11, center: bool = True, breaks: List[int] = None, center_breaks: bool = False) -> List[Images]:
     """Merge lines into images of maximum aspect ratio `aspect`.
 
     The resulting images will have the same width as the widest input line.
     
     The result may have aspect ratio larger than `aspect` if an input image already
     does; otherwise, they will all have height `max_width / aspect`.
 
     :param breaks: A list of forced page-breaks (e.g. page-turns).
     :param center: Center images horizontally.
     :param center_breaks: Space out lines on below-height pages.
     """
     width = max(i.width for i in images)
 
     todo = []
     output = []
     broken = []
     ## Split input images into groups of appropriate size
     for n, i in enumerate(images):
         height = sum((j.height for j in todo))
         if todo and (width/(height+i.height)) < aspect or n-1 in breaks:
             output.append(todo)
             broken.append(len(output))
             todo = []
 
         todo.append(i)
 
     output.append(todo)
 
     ## Merge them together
     for n, todo in enumerate(output, 1):
         height = sum((j.height for j in todo))
         outimage = Image.new("L", (width, max(height, int(width/aspect + 1))), color=255)
 
         if center and (center_breaks or n not in broken):
             padding = int((width/aspect - height)/(len(todo)+1))
         else:
             padding = 0
 
         cury = padding
         for t in todo:
             outimage.paste(t, (0, cury))
             cury += t.height + padding
 
         yield outimage
 
 DEFAULT_FONT = ImageFont.load_default(64)
 
 def number_images(images: List[Image], start: int = 0, in_place: bool = False, font = DEFAULT_FONT) -> List[Image]:
     res = []
     for n, i in enumerate(images, start):
         if not in_place:
             i = i.copy()
         d = ImageDraw.Draw(i)
         d.text((10, 10), str(n), font=font, fill=0)
 
         res.append(i)
 
     return res
 
 ## END Scenes
 
 if __name__ == "__main__":
     import argparse as ap
 
     parser = ap.ArgumentParser()
     subparsers = parser.add_subparsers(dest="command", required=True)
 
     find = subparsers.add_parser("find", help="Find scenes in video")
     find.add_argument("input", help="Input video to process")
     find.add_argument("-c", "--cutoff", type=float, default=0.03, help="Default scene cutoff for FFmpeg")
     find.add_argument("-l", "--line", action="store_true", help="Remove scrolling line")
     find.add_argument("-t", "--scene-time", type=float, default=2, help="Length of shortest scene")
     find.add_argument("-f", "--format", help="Output format", default="png")
     find.add_argument("outdir", help="Output directory")
 
     merge = subparsers.add_parser("merge", help="Merge scenes into a PDF")
     merge.add_argument("input", help="Input directory")
     merge.add_argument("-l", "--line-dist", default=10, type=int, help="Minimum distance between lines")
     merge.add_argument("-a", "--aspect", help="Desired aspect ratio (width/height)", type=float, default=8.5/11)
     merge.add_argument("-p", "--padding", help="Padding between cropped lines", type=int, default=50)
     merge.add_argument("-e", "--center", action="store_true", help="Horizontally center lines")
     merge.add_argument("-b", "--center-breaks", action="store_true", help="Vertically center breaks on a page")
     merge.add_argument("-s", "--savgol", type=float, default=0, help="Savgol threshold (0 to disable)")
     merge.add_argument("-o", "--otsu", action="store_true", help="Run Otsu thresholding")
     merge.add_argument("-f", "--format", help="Output format", default="png")
     merge.add_argument("-d", "--dpi", type=int, default=300, help="PDF DPI")
     merge.add_argument("-n", "--number", action="store_true", help="Number lines in output")
     merge.add_argument("--no-merge", action="store_true", help="Don't merge the output, save them to the output folder instead")
     merge.add_argument("output", help="Output file")
     merge.add_argument("forced_breaks", nargs="*", type=int, help="Forced page breaks")
 
     args = parser.parse_args()
 
     if args.command == "find":
         os.makedirs(args.outdir, exist_ok=True)
 
         scenes = find_scenes(args.input, args.cutoff, line=args.line, scene_cutoff=args.scene_time)
 
         fmt = "out%07d."+args.format
         for n, s in enumerate(scenes):
             s.save(os.path.join(args.outdir, fmt % n))
     elif args.command == "merge":
         if args.no_merge:
             os.makedirs(args.output, exist_ok=True)
 
         scenes = []
         for f in sorted(os.listdir(args.input)):
             if not f.endswith('.'+args.format):
                 continue
             scenes.append(Image.open(os.path.join(args.input, f)))
 
         new = []
         width = max(s.width for s in scenes)
         for s in scenes:
             s = crop_scene(s, width)
             for r in split_scene(s, args.line_dist):
                 new.append(pad_scene(r, args.padding))
 
         if args.number:
             number_images(new, in_place=True)
 
         with tempfile.TemporaryDirectory() as outdir:
             outfiles = []
             fmt = "out%07d."+args.format
 
             for n, done in enumerate(merge_lines(new, aspect=args.aspect, center=args.center, center_breaks=args.center_breaks, breaks=args.forced_breaks), 1):
                 name = os.path.join(outdir if not args.no_merge else args.output, fmt % n)
                 outfiles.append(name)
                 if args.otsu:
                     done = otsu(done.resize((done.width*2, done.height*2), resample=Image.Resampling.LANCZOS))
                 if args.savgol > 0:
                     done = do_savgol(done, args.savgol)
                 done.save(name)
 
             if not args.no_merge:
                 subp.check_call(["magick", "-density", str(args.dpi), *outfiles, args.output])