From 73aa653537feb01c4bd9989ec169740c49a9ef9a Mon Sep 17 00:00:00 2001
From: Kiran Sanjeevan <ksanjeevancabeza@gmail.com>
Date: Thu, 19 Jul 2018 20:42:38 -0700
Subject: [PATCH] Added gifs and things

---
 README.md        |  4 ++--
 net/netparams.py | 14 +++++++++-----
 net/utils.py     |  8 ++++++++
 yolov2.py        | 17 ++++++++++-------
 4 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 54c35aa..0181c24 100644
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ Original paper and github: [YOLO9000: Better, Faster, Stronger](https://arxiv.or
 - [h5py](http://www.h5py.org/)
 - [opencv](https://pypi.org/project/opencv-python/)
 - [python3.5](https://www.python.org/)
-
+- [moviepy](https://zulko.github.io/moviepy/) (optional, gifs)
 ### Simple use
 ---
 1. Download pretrained [model](https://drive.google.com/open?id=1khOgS8VD-paUD8KhjOLOzEXkzaXNAAMq) and place it in **dourflow/**.
@@ -192,7 +192,7 @@ Running:
 
 `python3 dourflow.py genp -c config.json`
 
-Will store your the bounding box priors wherever the path indicates in the config file under **config['config_path']['anchors']** with the prefix 'custom_' (so as to not overwrite accidentally).
+Will store the custom bounding box priors wherever the path indicates in the config file under **config['config_path']['anchors']** with the prefix 'custom_' (so as to not overwrite accidentally).
 
 ##### Tensorboard
 
diff --git a/net/netparams.py b/net/netparams.py
index a269054..8e2e578 100644
--- a/net/netparams.py
+++ b/net/netparams.py
@@ -42,6 +42,12 @@
         default='weights.h5')
 
 
+argparser.add_argument(
+        '--gif',
+        help='video output stored as gif also',
+        action='store_true')
+
+
 args = argparser.parse_args()
 
 
@@ -101,7 +107,7 @@ class YoloParams(object):
 
             PREDICT_IMAGE = action
         
-
+    #Paths 
     TRAIN_IMG_PATH = config['train']['image_folder'] 
     TRAIN_ANN_PATH = config['train']['annot_folder']
 
@@ -109,14 +115,12 @@ class YoloParams(object):
     VALIDATION_ANN_PATH = config['valid']['annot_folder']
     VALIDATION_OUT_PATH = config['valid']['pred_folder']
 
+    STORE_GIF = args.gif
+
     # Model    
-    #IN_MODEL = config['config_path']['in_model']
     IN_MODEL = args.model
-
     assert os.path.isfile(IN_MODEL), "Pass valid input keras model."
-
     OUT_MODEL_NAME = config['train']['out_model_name']
-    
     ARCH_FNAME = config['config_path']['arch_plotname']
 
     # Classes
diff --git a/net/utils.py b/net/utils.py
index 1ceae78..3bfc296 100755
--- a/net/utils.py
+++ b/net/utils.py
@@ -8,6 +8,8 @@
 import copy
 import cv2
 
+from moviepy.editor import VideoFileClip
+
 
 
 def mkdir_p(path):
@@ -230,6 +232,12 @@ def handle_empty_indexing(arr, idx):
     return []
 
 
+def generate_gif(filename):
+    outname = filename.split('.')[-2] + '.gif'
+    VideoFileClip(filename).write_gif(
+        outname,fps=20, program='ffmpeg', fuzz=3)
+    print('\n')
+
 
 if __name__ == '__main__':
 
diff --git a/yolov2.py b/yolov2.py
index 98a6f38..7c7d2c0 100644
--- a/yolov2.py
+++ b/yolov2.py
@@ -12,7 +12,7 @@
 from keras.optimizers import SGD, Adam, RMSprop
 
 from net.utils import parse_annotation, mkdir_p, \
-setup_logging, draw_boxes
+setup_logging, draw_boxes, generate_gif
 
 from net.netparams import YoloParams
 from net.netloss import YoloLoss
@@ -23,8 +23,6 @@
 
 
 
-
-
 class YoloV2(object):
 
 
@@ -107,9 +105,9 @@ def _video_params(self, name):
     def video_inference(self, filename):
 
         cap, size, video_len, fps = self._video_params(filename)
-
+        outname = filename.split('.')[0]+"_pred.mp4"
         fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
-        writer = cv2.VideoWriter(filename.split('.')[0]+"_pred.mp4", fourcc, fps, size)
+        writer = cv2.VideoWriter(outname, fourcc, fps, size)
         
         for i in tqdm(range(video_len)):
 
@@ -127,13 +125,16 @@ def video_inference(self, filename):
         writer.release()
         cv2.destroyAllWindows()
 
+        if YoloParams.STORE_GIF: generate_gif(outname)
+
+
 
     def cam_inference(self, fname):
 
         cap, size, _, fps = self._video_params(0)
 
         fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
-        if fname: writer = cv2.VideoWriter("out.mp4", fourcc, fps, size)
+        if fname: writer = cv2.VideoWriter(fname, fourcc, 40, size)
 
         while(cap.isOpened()):
 
@@ -150,7 +151,9 @@ def cam_inference(self, fname):
                     break
 
         cap.release()
-        if fname: writer.release()
+        if fname: 
+            writer.release()
+            if YoloParams.STORE_GIF: generate_gif(fname)
         cv2.destroyAllWindows()