Add main CLI function and library function for getting segment times

main.py

@@ -0,0 +1,21 @@
+# Import needed modules
+import video2gif
+import optparse
+import json
+from moviepy.editor import VideoFileClip
+
+def parser():
+    parser = optparse.OptionParser()
+    parser.add_option("-s", "--source", default="./videos/FrG4TEcSuRg.mp4", help="Which video to process")
+    parser.add_option("-d", "--duration", default=3, help="Duration of the segments", type="int")
+    parser.add_option("-t", "--top", default=5, help="How many top segments to get", type="int")
+    parser.add_option("-b", "--bottom", default=0, help="How many bottom segments to get", type="int")
+    return parser.parse_args()
+
+def main():
+    args, opts = parser()
+    scored_segments = get_scored_segments(args.video, args.duration, args.top, args.bottom)
+    print(json.dumps(scored_segments, indent=4))
+
+if __name__ == "__main__":
+    main()

video2gif/__init__.py

@@ -175,3 +175,62 @@ def generate_gifs(out_dir, segment2scores, video, video_id, top_k=6, bottom_k=0)
         nr -= 1
 
     return good_gifs,bad_gifs
+
+
+def generate_gif_times(video, segment2scores, top_k=6, bottom_k=0):
+    '''
+    @param out_dir: directory where the GIFs are written to
+    @param segment2scores: a dict with segments (start frame, end frame) as keys and the segment score as value
+    @param video: a VideoFileClip object
+    @param video_id: the identifier of the video (used for naming the GIFs)
+    @return:
+    '''
+    segment2scores = segment2scores.copy()
+    print("found segscors", len(segment2scores))
+
+    nr=0
+    top_k=min(top_k, len(segment2scores))
+    good_gifs=[]
+    for segment in sorted(segment2scores, key=lambda x: -segment2scores.get(x))[0:top_k]:
+        segment_times = dict(start=segment[0]/float(video.fps), end=segment[1]/float(video.fps))
+        good_gifs.append(segment_times)
+        nr += 1
+
+    bottom_k=min(bottom_k, len(segment2scores))
+    bad_gifs=[]
+    nr=len(segment2scores)
+    for segment in sorted(segment2scores, key=segment2scores.get)[0:bottom_k]:
+        segment_times = dict(start=segment[0]/float(video.fps), end=segment[1]/float(video.fps))
+        bad_gifs.append(segment_times)
+        nr -= 1
+
+    return good_gifs, bad_gifs
+
+
+# Define function
+def get_scored_segments( video_path, duration = 3, top_k = 5, bottom_k = 0 ):
+    '''
+    @param video_path: video to run gif segment scoring on
+    @param duration: duration of segments
+    @param top_k: count of top gifs to return
+    @param bottom_k: count of bottom gifs to return
+    @return: @object(good, bad)
+    '''
+
+    # Get scoring function
+    score_function = video2gif.get_prediction_function()
+
+    # Take the example video
+    video = VideoFileClip(video_path)
+
+    # Build the segments
+    segments = [(start, int(start+video.fps*duration)) for start in range(0,int(video.duration*video.fps),int(video.fps*duration))]
+
+    # Score the segments
+    scores= get_scores(score_function, segments, video, stride=8)
+
+    # Generate GIFs from the top scoring segments
+    good_gifs, bad_gifs = generate_gif_times(video, scores, top_k, bottom_k)
+
+    # Return the good and bad gifs
+    return dict(good=good_gifs, bad=bad_gifs)