gigant/slideshow_extraction

1BATCH_SIZE = 642DOWNSAMPLE = 243FOLDER_PATH = "."4 5import phash_jax6import jax.numpy as jnp7import matplotlib.pyplot as plt8from PIL import Image9import statistics10from decord import VideoReader11from decord import cpu12import gradio13 14def binary_array_to_hex(arr):15	"""16	Function to make a hex string out of a binary array.17	"""18	bit_string = ''.join(str(b) for b in 1 * arr.flatten())19	width = int(jnp.ceil(len(bit_string) / 4))20	return '{:0>{width}x}'.format(int(bit_string, 2), width=width)21 22def compute_batch_hashes(vid_path):23  kwargs={"width": 64, "height":64}24  vr = VideoReader(vid_path, ctx=cpu(0), **kwargs)25  hashes = []26  h_prev = None27  batch = []28  for i in range(0, len(vr), DOWNSAMPLE * BATCH_SIZE):29      print(f"batch_{i}")30      ids = [id for id in range(i, min(i + DOWNSAMPLE * BATCH_SIZE, len(vr)), DOWNSAMPLE)]31      vr.seek(0)32      batch = jnp.array(vr.get_batch(ids).asnumpy())33      batch_h =  phash_jax.batch_phash(batch)34      for i in range(len(ids)):35        h = batch_h[i]36        if h_prev == None:37          h_prev=h38        hashes.append({"frame_id":ids[i], "hash": binary_array_to_hex(h), "distance": int(phash_jax.hash_dist(h, h_prev))})39        h_prev = h40  return gradio.update(value=hashes, visible=False)41 42def plot_hash_distance(hashes, threshold):43  fig = plt.figure()44  ids = [h["frame_id"] for h in hashes]45  distances = [h["distance"] for h in hashes]46  plt.plot(ids, distances, ".")47  plt.plot(ids, [threshold]* len(ids), "r-")48  return fig49 50def compute_threshold(hashes):51  min_length = 24 * 352  ids = [h["frame_id"] for h in hashes]53  distances = [h["distance"] for h in hashes]54  thrs_ = sorted(list(set(distances)),reverse=True)55  best = thrs_[0] - 156  for threshold in thrs_[1:]:57    durations = []58    i_start=059    for i, h in enumerate(hashes):60      if h["distance"] > threshold and hashes[i-1]["frame_id"] -  hashes[i_start]["frame_id"] > min_length:61        durations.append(hashes[i-1]["frame_id"] - hashes[i_start]["frame_id"])62        i_start=i63    if len(durations) < (len(hashes) * DOWNSAMPLE / 24) / 20:64      best = threshold65  return best66 67def get_slides(vid_path, hashes, threshold):68    min_length = 24 * 1.569    vr = VideoReader(vid_path, ctx=cpu(0))70    slideshow = []71    i_start = 072    for i, h in enumerate(hashes):73        if h["distance"] > threshold and hashes[i-1]["frame_id"] -  hashes[i_start]["frame_id"] > min_length:74            path=f'{FOLDER_PATH}/{vid_path.split("/")[-1].split(".")[0]}_{i_start}_{i-1}.png'75            Image.fromarray(vr[hashes[i-1]["frame_id"]].asnumpy()).save(path)76            slideshow.append({"slide": path, "start": i_start, "end": i-1})77            i_start=i78    path=f'{FOLDER_PATH}/{vid_path.split("/")[-1].split(".")[0]}_{i_start}_{len(vr)-1}.png'79    Image.fromarray(vr[-1].asnumpy()).save(path)80    slideshow.append({"slide": path, "start": i_start, "end": len(vr)-1})81    return [s["slide"] for s in slideshow]82 83def trigger_plots(f2f_distance_plot, hashes, threshold):84  # if not hist_plot.get_config()["visible"] and len(hashes.get_config()["value"]) > 0 :85  return gradio.update(value=plot_hash_distance(hashes, threshold))86 87def set_visible():88  return gradio.update(visible=True)89  90demo = gradio.Blocks()91 92with demo:93    with gradio.Row():94        with gradio.Column():95          with gradio.Row():96            vid=gradio.Video(mirror_webcam=False)97          with gradio.Row():98            btn_vid_proc = gradio.Button("Compute hashes")99          with gradio.Row():100            hist_plot = gradio.Plot(label="Frame to frame hash distance histogram", visible=False)101        with gradio.Column():102            hashes = gradio.JSON()103            with gradio.Column(visible=False) as result_row:104              btn_plot = gradio.Button("Plot & compute optimal threshold")105              threshold = gradio.Slider(minimum=1, maximum=30, value=5, label="Threshold")106              f2f_distance_plot = gradio.Plot(label="Frame to frame hash distance")107              btn_slides = gradio.Button("Extract Slides")108    with gradio.Row():109      slideshow = gradio.Gallery(label="Extracted slides", columns=[6], rows=[1])110      # slideshow.style(grid=6)111    btn_vid_proc.click(fn=compute_batch_hashes, inputs=[vid], outputs=[hashes])112    hashes.change(fn=set_visible, inputs=[], outputs=[result_row])113    btn_plot.click(fn=compute_threshold, inputs=[hashes], outputs=[threshold])114    btn_plot.click(fn=trigger_plots, inputs=[f2f_distance_plot, hashes, threshold], outputs=[f2f_distance_plot])115    threshold.change(fn=plot_hash_distance, inputs=[hashes, threshold], outputs=f2f_distance_plot)116    btn_slides.click(fn=get_slides, inputs=[vid, hashes, threshold], outputs=[slideshow])117 118demo.queue().launch()