From 2eb793c992eed3a5f6ab397b75c62185af8925d6 Mon Sep 17 00:00:00 2001
From: Kensho Hara <kensho.hara@aist.go.jp>
Date: Mon, 25 Sep 2017 16:12:36 +0900
Subject: [PATCH] add feature mode

---
 classify.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/classify.py b/classify.py
index 371247a..78550f8 100644
--- a/classify.py
+++ b/classify.py
@@ -6,6 +6,8 @@
 from temporal_transforms import LoopPadding
 
 def classify_video(video_dir, video_name, class_names, model, opt):
+    assert opt.mode in ['score', 'feature']
+
     spatial_transform = Compose([Scale(opt.sample_size),
                                  CenterCrop(opt.sample_size),
                                  ToTensor(),
@@ -37,9 +39,14 @@ def classify_video(video_dir, video_name, class_names, model, opt):
     for i in range(video_outputs.size(0)):
         clip_results = {
             'segment': video_segments[i].tolist(),
-            'label': class_names[max_indices[i]],
-            'scores': video_outputs[i].tolist()
         }
+
+        if opt.mode == 'score':
+            clip_results['label'] = class_names[max_indices[i]]
+            clip_results['scores'] = video_outputs[i].tolist()
+        elif opt.mode == 'feature':
+            clip_results['features'] = video_outputs[i].tolist()
+
         results['clips'].append(clip_results)
 
     return results