Skip to content

Custom import

To integrate custom import plugin, you will need to import our Supervisely python library and create simple class with a few methods.

In this example we will implement a new plugin to import video formats into Supervisely.

Hint

Look for up-to-date sources on GitHub.

In order to make import plugin, developer needs to code a few stages:

  1. Read paths and settings

  2. Get list of input files / datasets

  3. Generate output project folders and meta.json.

  4. Process and write images and annotations to output project according to datasets

Installing environment

Assuming you have Linux based OS and Docker (nvidia-docker) already

Step 1. Create project folder

    $ mkdir import_test
    $ cd import_test

Step 2. Clone Supervisely library from github

    $ git clone https://github.com/supervisely/supervisely.git

Step 3. Create src folder and create main.py script test file

    $ mkdir src
    $ cd src
    $ echo "print('Hello')" > main.py
    $ cd ..

Step 4. Create Dockerfile

    $ nano Dockerfile 

and copy-paste the following content:

FROM supervisely/base-py:latest
ENV PYTHONPATH /workdir:/workdir/supervisely_lib/worker_proto:$PYTHONPATH

WORKDIR /workdir/src

COPY ./supervisely/supervisely_lib /workdir/supervisely_lib
COPY src /workdir/src

ENTRYPOINT ["sh", "-c", "python -u main.py"]

Step 5. Build Docker image

    $ docker build -t "supervisely/custom_video_import:ver_01" .

Step 6. Run Docker container

    $ nvidia-docker run --rm -t -i "supervisely/custom_video_import:ver_01"

You should see Hello text in your shell.

Import configuration

Before coding let's take a look at default configuration for this plugin:

{
    "options": {
        "skip_frame": 25
    },

    "res_names": {
        "project": "test_project"
    }
}

Code

All code we will write in main.py file, created above.

At first - import all required libs:

import os
import cv2

import supervisely_lib as sly
from supervisely_lib import logger

Then define VideoImporter class and write constructor:

class VideoImporter:
    def __init__(self):
        # Gets all path's we need
        task_paths = sly.DtlPaths()
        self.in_dir = task_paths.data_dir
        self.out_dir = task_paths.results_dir

        # Load task settings into variable
        self.settings = sly.json_load(task_paths.settings_path)

        # Extract task options that we need
        self.skip_frame = int(self.settings['options'].get('skip_frame', 25))

Main process method will get video files names and put into output project dataset structure for each video file.

    def process(self):
        # Getting video files list, filtered by allowed file extensions
        video_files_list = sly.ImportVideoLister.list_videos(self.in_dir)

        logger.info('Read videos from source:', extra={'count': len(video_files_list)})

        # Generate output project items info (list if items)
        out_pr = sly.ProjectStructure(self.settings['res_names']['project'])
        for video_file in video_files_list:
            file_name, file_ext = os.path.splitext(video_file)
            src_video_file_path = os.path.join(self.in_dir, video_file)
            meta_data = {
                'src_video_file_path': src_video_file_path,
                'image_ext': file_ext  # - Required field
            }
            out_pr.add_item(file_name, file_name, meta_data)

        # Generate output project paths and save project-meta information
        out_pr_fs = sly.ProjectFS(self.out_dir, out_pr)
        out_pr_fs.make_dirs()
        res_meta = sly.ProjectMeta()
        res_meta.to_dir(out_pr_fs.project_path)

        # Process samples (video files)
        for sample_info in out_pr_fs:
            self.process_video_file(sample_info, out_pr)

Method process_video_file does the following steps:

  1. Open video capture

  2. Read current frame

  3. Save frame as an image file

  4. Save empty annotation related to that image

  5. Repeat for each frame

Consider the following code:

    def process_video_file(self, sample_info, out_pr):
        try:
            video_capture = cv2.VideoCapture(sample_info.ia_data["src_video_file_path"])

            # Get total frames count in current video file
            length = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))

            # We are to report progress of task over sly.ProgressCounter if we want to observe the progress in web panel.
            progress = sly.progress_counter_import(out_pr.name, length)

            # Read frames from video video capture
            while video_capture.isOpened():
                ret, image = video_capture.read()

                # If frames is out, we need stop process
                if ret is not True:
                    break

                # Get current frame number
                frame_id = int(video_capture.get(cv2.CAP_PROP_POS_FRAMES))

                # Skip few frames if need
                if frame_id % self.skip_frame == 0:

                    # Generate names for files
                    frame_name = "{0}_frame_{1}".format(sample_info.image_name, str(frame_id).zfill(7))
                    img_frame_path = os.path.join(os.path.dirname(sample_info.img_path), frame_name + '.jpg')
                    json_frame_path = os.path.join(os.path.dirname(sample_info.ann_path), frame_name + '.json')

                    # Write image
                    sly.write_image(img_frame_path, image)

                    # Write empty annotation
                    imsize_wh = image.shape[:2]
                    ann = sly.Annotation.new_with_objects(imsize_wh, [])
                    sly.json_dump(ann.pack(), json_frame_path)

                    progress.iter_done_report()
            video_capture.release()

        except Exception as e:
            exc_str = str(e)
            logger.warn('Input sample skipped due to error: {}'.format(exc_str), exc_info=True, extra={
                'exc_str': exc_str,
                'dataset_name': sample_info.ds_name,
                'image_name': sample_info.image_name,
            })

Finally, you may wrap your main function as shown below:

def main():
    importer = VideoImporter()
    importer.process()


if __name__ == '__main__':
    sly.main_wrapper('VIDEO_ONLY_IMPORT', main)

Supervisely main_wrapper function can handle root exceptions and log it to the UI.

Deploy

Build Docker image again with updated src/main.py script:

    $ docker build -t "supervisely/custom_video_import:ver_01" .

Push Docker image:

    $ docker push -t "supervisely/custom_video_import:ver_01"

Create new import plugin at Plugins page and choose Docker image that you have pushed above.

Don't forget to set default config.