diff --git a/main.py b/main.py new file mode 100644 index 0000000..593588e --- /dev/null +++ b/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() diff --git a/video2gif/__init__.py b/video2gif/__init__.py index d1aa19d..881dcf4 100644 --- a/video2gif/__init__.py +++ b/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)