def extract_features()

in tools/scripts/features/extract_features_vmb.py [0:0]

• 43 lines of code
• 13 McCabe index (conditional complexity)

    def extract_features(self):
        image_dir = self.args.image_dir

        if os.path.isfile(image_dir):
            features, infos = self.get_detectron_features([image_dir])
            self._save_feature(image_dir, features[0], infos[0])
        else:
            files = glob.glob(os.path.join(image_dir, "*.png"))
            files.extend(glob.glob(os.path.join(image_dir, "*.jpg")))
            files.extend(glob.glob(os.path.join(image_dir, "*.jpeg")))
            files = {f: 1 for f in files}
            exclude = {}

            if os.path.exists(self.args.exclude_list):
                with open(self.args.exclude_list) as f:
                    lines = f.readlines()
                    for line in lines:
                        exclude[
                            line.strip("\n").split(os.path.sep)[-1].split(".")[0]
                        ] = 1
            output_files = glob.glob(os.path.join(self.args.output_folder, "*.npy"))
            output_dict = {}
            for f in output_files:
                file_name = f.split(os.path.sep)[-1].split(".")[0]
                output_dict[file_name] = 1

            for f in list(files.keys()):
                file_name = f.split(os.path.sep)[-1].split(".")[0]
                if file_name in output_dict or file_name in exclude:
                    files.pop(f)

            files = list(files.keys())

            finished = 0

            end_index = self.args.end_index
            if end_index is None:
                end_index = len(files)
            start_index = self.args.start_index

            total = len(files[start_index:end_index])

            for chunk in self._chunks(
                files[start_index:end_index], self.args.batch_size
            ):
                features, infos = self.get_detectron_features(chunk)
                for idx, file_name in enumerate(chunk):
                    self._save_feature(file_name, features[idx], infos[idx])
                finished += len(chunk)

                if finished % 200 == 0:
                    print(f"Processed {finished}/{total}")