add feature extractor

Author: joonson

SyncNetInstance.py

@@ -36,7 +36,7 @@ class SyncNetInstance(torch.nn.Module):
     def __init__(self, dropout = 0, num_layers_in_fc_layers = 1024):
         super(SyncNetInstance, self).__init__();
 
-        self.__S__ = S(num_layers_in_fc_layers = num_layers_in_fc_layers);
+        self.__S__ = S(num_layers_in_fc_layers = num_layers_in_fc_layers).cuda();
 
     def evaluate(self, opt, videofile):
 
@@ -139,6 +139,57 @@ def evaluate(self, opt, videofile):
         dists_npy = numpy.array([ dist.numpy() for dist in dists ])
         return offset.numpy(), conf.numpy(), dists_npy
 
+    def extract_feature(self, opt, videofile):
+
+        self.__S__.eval();
+        
+        # ========== ==========
+        # Load video 
+        # ========== ==========
+        cap = cv2.VideoCapture(videofile)
+
+        frame_num = 1;
+        images = []
+        while frame_num:
+            frame_num += 1
+            ret, image = cap.read()
+            if ret == 0:
+                break
+
+            image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            images.append(image_np)
+
+        im = numpy.stack(images,axis=3)
+        im = numpy.expand_dims(im,axis=0)
+        im = numpy.transpose(im,(0,3,4,1,2))
+
+        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())
+        
+        # ========== ==========
+        # Generate video feats
+        # ========== ==========
+
+        lastframe = len(images)-4
+        im_feat = []
+
+        tS = time.time()
+        for i in range(0,lastframe,opt.batch_size):
+            
+            im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
+            im_in = torch.cat(im_batch,0)
+            im_out  = self.__S__.forward_lip(im_in.cuda());
+            im_feat.append(im_out.data.cpu())
+
+        im_feat = torch.cat(im_feat,0)
+
+        # ========== ==========
+        # Compute offset
+        # ========== ==========
+            
+        print('Compute time %.3f sec.' % (time.time()-tS))
+
+        return im_feat
+
 
     def loadParameters(self, path):
         loaded_state = torch.load(path, map_location=lambda storage, loc: storage);

SyncNetModel.py

@@ -89,12 +89,6 @@ def __init__(self, num_layers_in_fc_layers = 1024):
             nn.BatchNorm3d(2048),
             nn.ReLU(inplace=True),
         );
-        
-        self.netcnnaud = self.netcnnaud.cuda();
-        self.netcnnlip = self.netcnnlip.cuda();
-        self.netfcaud  = self.netfcaud.cuda();
-        self.netfclip = self.netfclip.cuda();
-        
 
     def forward_aud(self, x):
 
@@ -110,4 +104,10 @@ def forward_lip(self, x):
         mid = mid.view((mid.size()[0], -1)); # N x (ch x 24)
         out = self.netfclip(mid);
 
+        return out;
+
+    def forward_lipfeat(self, x):
+
+        out = self.netcnnlip(x);
+
         return out;

demo_feature.py

@@ -0,0 +1,32 @@
+#!/usr/bin/python
+#-*- coding: utf-8 -*-
+
+import time, pdb, argparse, subprocess
+
+from SyncNetInstance import *
+
+# ==================== LOAD PARAMS ====================
+
+
+parser = argparse.ArgumentParser(description = "SyncNet");
+
+parser.add_argument('--initial_model', type=str, default="data/syncnetl2.model", help='');
+parser.add_argument('--batch_size', type=int, default='20', help='');
+parser.add_argument('--vshift', type=int, default='15', help='');
+parser.add_argument('--videofile', type=str, default="data/example.avi", help='');
+parser.add_argument('--tmp_dir', type=str, default="data", help='');
+parser.add_argument('--save_as', type=str, default="data/features.pt", help='');
+
+opt = parser.parse_args();
+
+
+# ==================== RUN EVALUATION ====================
+
+s = SyncNetInstance();
+
+s.loadParameters(opt.initial_model);
+print("Model %s loaded."%opt.initial_model);
+
+feats = s.extract_feature(opt, videofile=opt.videofile)
+
+torch.save(feats, opt.save_as)

demo_syncnet.py

@@ -13,8 +13,8 @@
 parser.add_argument('--initial_model', type=str, default="data/syncnetl2.model", help='');
 parser.add_argument('--batch_size', type=int, default='20', help='');
 parser.add_argument('--vshift', type=int, default='15', help='');
-parser.add_argument('--videofile', type=str, default="", help='');
-parser.add_argument('--tmp_dir', type=str, default="~", help='');
+parser.add_argument('--videofile', type=str, default="data/example.avi", help='');
+parser.add_argument('--tmp_dir', type=str, default="data", help='');
 
 opt = parser.parse_args();