add auto download fucntion

.gitignore

@@ -7,6 +7,5 @@ tests
 dist
 *.pyc
 __pycache__
-pretrained
-cen_patches_*_of.dat
+exts
 tmp.csv

MANIFEST.in

@@ -1,5 +1 @@
-recursive-include src/MSA_FET/exts *
-recursive-include src/MSA_FET/example_configs *.json
-
-recursive-exclude src/MSA_FET/exts/pretrained *
-recursive-exclude src/MSA_FET/exts/OpenFace/model/patch_experts *.dat
+recursive-include src/MSA_FET/example_configs *.json

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = MMSA-FET
-version = 0.2.8
+version = 0.2.9
 author = THUIAR
 maintainer = FlameSky
 description = A Tool for extracting multimodal features from videos.
@@ -36,6 +36,7 @@ install_requires =
     numpy >= 1.20.3
     pandas >= 1.2.5
     tqdm >= 4.62.2
+    gdown >= 4.4.0
 
 [options.packages.find]
 where = src

src/MSA_FET/ASD/model/faceDetector/s3fd/__init__.py

@@ -1,5 +1,3 @@
-import os
-import subprocess
 from pathlib import Path
 
 import cv2
@@ -10,24 +8,23 @@
 from .nets import S3FDNet
 
 PATH_WEIGHT = Path(__file__).parent.parent.parent.parent.parent / 'exts' / 'pretrained' / 'sfd_face.pth'
-if os.path.isfile(PATH_WEIGHT) == False:
-    Link = "1KafnHz7ccT-3IyddBsL5yi2xGtxAKypt"
-    cmd = "gdown --id %s -O %s"%(Link, PATH_WEIGHT)
-    subprocess.call(cmd, shell=True, stdout=None)
 img_mean = np.array([104., 117., 123.])[:, np.newaxis, np.newaxis].astype('float32')
 
 
 class S3FD():
 
     def __init__(self, device='cuda'):
 
+        if not PATH_WEIGHT.exists():
+            print("Model 'sdf_face.pth' not found. Please run `python -m MSA_FET install` to download pretrained files.")
+
         # tstamp = time.time()
         self.device = device
 
         # print('[S3FD] loading with', self.device)
         self.net = S3FDNet(device=self.device).to(self.device)
-        PATH = os.path.join(os.getcwd(), PATH_WEIGHT)
-        state_dict = torch.load(PATH, map_location=self.device)
+        # PATH = os.path.join(os.getcwd(), PATH_WEIGHT)
+        state_dict = torch.load(PATH_WEIGHT, map_location=self.device)
         self.net.load_state_dict(state_dict)
         self.net.eval()
         # print('[S3FD] finished loading (%.4f sec)' % (time.time() - tstamp))

src/MSA_FET/__main__.py

@@ -2,74 +2,97 @@
 from pathlib import Path
 
 from .main import FeatureExtractionTool
+from .install import download_missing, force_redownload
 
 
 def parse_args():
+    # main parser
     parser = argparse.ArgumentParser()
-    parser.add_argument('-i', '--input', type=str, required=True,
-                        help="Input video file in file mode, or dataset dir in dataset mode.")
-    parser.add_argument('-d', '--dataset-mode', action='store_true',
-                        help="Switch from file mode to dataset mode if specified.")
-    parser.add_argument('-c', '--config-file', type=str, required=True,
-                        help="Path to config file.")
-    parser.add_argument('-o', '--output', type=str, required=True, # can only write to file in commandline
-                        help="Path to output pkl file.")
-    parser.add_argument('-t', '--text-file', type=str, default=None,
-                        help="File containing transcriptions of the video. Required for extracting text features for single video.")
-    parser.add_argument('-n', '--num-workers', type=int, default=4,
-                        help="Number of workers for data loading in dataset mode.")
+    subparsers = parser.add_subparsers(dest="command")
     parser.add_argument('-v', '--verbose', action='store_true',
                         help="Print more information to stdout.")
     parser.add_argument('-q', '--quiet', action='store_true',
                         help="Print only errors to stdout.")
-    parser.add_argument('--return_type', type=str, default='np', choices=['np', 'pt', 'df'],
-                        help="Return type of the tool.")
     parser.add_argument('--tmp-dir', type=str, default=Path.home() / '.MMSA-FET/tmp',
                         help="Temporary directory for storing intermediate results. Default: '~/.MSA-FET/tmp'")
     parser.add_argument('--log-dir', type=str, default=Path.home() / '.MMSA-FET/log',
                         help="Log file directory. Default: '~/.MSA-FET/log'")
-    parser.add_argument('--batch-size', type=int, default=4,
-                        help="Batch size for dataset mode. Default: 4")
-    parser.add_argument()
+
+    # sub-parser for installation
+    parser_install = subparsers.add_parser("install", help="Download models required to run MSA-FET")
+    parser_install.add_argument('-p', '--proxy', type=str, default=None,
+                                help="Proxy for downloading. e.g. socks5://127.0.0.1:8080")
+    parser_install.add_argument('-f', '--force', action='store_true', help="Force re-download models.")
+
+    # sub-parser for dataset mode
+    parser_single = subparsers.add_parser("run_single", help="Run feature extraction on a single video")
+    parser_single.add_argument('-i', '--input', type=str, required=True,
+                                help="Input video file in file mode, or dataset dir in dataset mode.")
+    parser_single.add_argument('-c', '--config-file', type=str, required=True,
+                                help="Path to config file.")
+    parser_single.add_argument('-o', '--output', type=str, required=True,
+                                help="Path to output pkl file.")
+    parser_single.add_argument('-t', '--text-file', type=str, default=None,
+                                help="File containing transcriptions of the video. Required when extracting text features.")
+    parser_single.add_argument('--return_type', type=str, default='np', choices=['np', 'pt', 'df'],
+                                help="Return type of the tool.")
+
+
+    parser_dataset = subparsers.add_parser("run_dataset", help="Run feature extraction on a dataset")
+    parser_dataset.add_argument('-i', '--input', type=str, required=True,
+                                help="Input dataset dir.")
+    parser_dataset.add_argument('-c', '--config-file', type=str, required=True,
+                                help="Path to config file.")
+    parser_dataset.add_argument('-o', '--output', type=str, required=True,
+                                help="Path to output pkl file.")
+    parser_dataset.add_argument('-n', '--num-workers', type=int, default=4,
+                                help="Num of dataloader workers.")
+    parser_dataset.add_argument('--batch-size', type=int, default=4,
+                                help="Batch size. Default: 4")
+    parser_dataset.add_argument('--return_type', type=str, default='np', choices=['np', 'pt', 'df'],
+                                help="Return type of the tool.")
+
     return parser.parse_args()
 
 
 if __name__ == '__main__':
     args = parse_args()
 
-    verbose = 1
-    if args.verbose:
-        verbose = 2
-    if args.quiet:
-        verbose = 0
-
-    # if args.tmp_dir == '':
-    #     args.tmp_dir = Path.home() / '.MMSA-FET/tmp'
-    # if args.log_dir == '':
-    #     args.log_dir = Path.home() / '.MMSA-FET/log'
-
-    # if args.text_file == '':
-    #     args.text_file = None
-
-    fet = FeatureExtractionTool(
-        config=args.config_file,
-        tmp_dir=args.tmp_dir,
-        log_dir=args.log_dir,
-        verbose=verbose
-    )
-
-    if args.dataset_mode:
-        fet.run_dataset(
-            dataset_dir=args.input,
-            out_file=args.output,
-            num_workers=args.num_workers,
-            return_type=args.return_type,
-            batch_size=args.batch_size
-        )
+    if args.command == 'install':
+        if args.force:
+            force_redownload(args.proxy)
+        else:
+            download_missing(args.proxy)
+    elif args.command is None:
+        print("Please specify a command. Use -h to see all available commands.")
     else:
-        fet.run_single(
-            in_file=args.input,
-            out_file=args.output,
-            text_file=args.text_file,
-            return_type=args.return_type
+        verbose = 1
+        if args.verbose:
+            verbose = 2
+        if args.quiet:
+            verbose = 0
+
+        fet = FeatureExtractionTool(
+            config=args.config_file,
+            tmp_dir=args.tmp_dir,
+            log_dir=args.log_dir,
+            verbose=verbose
         )
+
+        if args.command == 'run_single':
+            fet.run_single(
+                in_file=args.input,
+                out_file=args.output,
+                text_file=args.text_file,
+                return_type=args.return_type
+            )
+        elif args.command == 'run_dataset':
+            fet.run_dataset(
+                dataset_dir=args.input,
+                out_file=args.output,
+                num_workers=args.num_workers,
+                return_type=args.return_type,
+                batch_size=args.batch_size
+            )
+        else:
+            print("Unknown command: {}".format(args.command))