From 4298166403bef25fcc90edc5a3e527ca619ea45b Mon Sep 17 00:00:00 2001 From: zliang7 Date: Fri, 31 May 2019 04:31:44 +0800 Subject: [PATCH] Support frame selection when create from video (#437) --- CHANGELOG.md | 1 + .../static/dashboard/js/dashboard.js | 53 +++++++++++++++++++ .../templates/dashboard/dashboard.html | 27 ++++++++++ cvat/apps/engine/annotation.py | 5 +- .../engine/migrations/0019_frame_selection.py | 40 ++++++++++++++ cvat/apps/engine/models.py | 11 ++-- cvat/apps/engine/serializers.py | 15 +++++- cvat/apps/engine/task.py | 27 +++++++--- 8 files changed, 168 insertions(+), 11 deletions(-) create mode 100644 cvat/apps/engine/migrations/0019_frame_selection.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 69b313da..d40ae678 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Installation guide - Linear interpolation for a single point +- Video frame filter ### Changed - Outside and keyframe buttons in the side panel for all interpolation shapes (they were only for boxes before) diff --git a/cvat/apps/dashboard/static/dashboard/js/dashboard.js b/cvat/apps/dashboard/static/dashboard/js/dashboard.js index 8bc302b6..8a0840c2 100644 --- a/cvat/apps/dashboard/static/dashboard/js/dashboard.js +++ b/cvat/apps/dashboard/static/dashboard/js/dashboard.js @@ -464,6 +464,10 @@ class DashboardView { return (overlapSize >= 0 && overlapSize <= segmentSize - 1); } + function validateStopFrame(stopFrame, startFrame) { + return !customStopFrame.prop('checked') || stopFrame >= startFrame; + } + function requestCreatingStatus(tid, onUpdateStatus, onSuccess, onError) { function checkCallback() { $.get(`/api/v1/tasks/${tid}/status`).done((data) => { @@ -516,6 +520,12 @@ class DashboardView { const customOverlapSize = $('#dashboardCustomOverlap'); const imageQualityInput = $('#dashboardImageQuality'); const customCompressQuality = $('#dashboardCustomQuality'); + const startFrameInput = $('#dashboardStartFrame'); + const customStartFrame = $('#dashboardCustomStart'); + const stopFrameInput = $('#dashboardStopFrame'); + const customStopFrame = $('#dashboardCustomStop'); + const frameFilterInput = $('#dashboardFrameFilter'); + const customFrameFilter = $('#dashboardCustomFilter'); const taskMessage = $('#dashboardCreateTaskMessage'); const submitCreate = $('#dashboardSubmitTask'); @@ -529,6 +539,9 @@ class DashboardView { let segmentSize = 5000; let overlapSize = 0; let compressQuality = 50; + let startFrame = 0; + let stopFrame = 0; + let frameFilter = ''; let files = []; dashboardCreateTaskButton.on('click', () => { @@ -612,6 +625,9 @@ class DashboardView { customSegmentSize.on('change', (e) => segmentSizeInput.prop('disabled', !e.target.checked)); customOverlapSize.on('change', (e) => overlapSizeInput.prop('disabled', !e.target.checked)); customCompressQuality.on('change', (e) => imageQualityInput.prop('disabled', !e.target.checked)); + customStartFrame.on('change', (e) => startFrameInput.prop('disabled', !e.target.checked)); + customStopFrame.on('change', (e) => stopFrameInput.prop('disabled', !e.target.checked)); + customFrameFilter.on('change', (e) => frameFilterInput.prop('disabled', !e.target.checked)); segmentSizeInput.on('change', () => { const value = Math.clamp( @@ -646,6 +662,28 @@ class DashboardView { compressQuality = value; }); + startFrameInput.on('change', function() { + let value = Math.max( + +startFrameInput.prop('value'), + +startFrameInput.prop('min') + ); + + startFrameInput.prop('value', value); + startFrame = value; + }); + stopFrameInput.on('change', function() { + let value = Math.max( + +stopFrameInput.prop('value'), + +stopFrameInput.prop('min') + ); + + stopFrameInput.prop('value', value); + stopFrame = value; + }); + frameFilterInput.on('change', function() { + frameFilter = frameFilterInput.prop('value'); + }); + submitCreate.on('click', () => { if (!validateName(name)) { taskMessage.css('color', 'red'); @@ -677,6 +715,12 @@ class DashboardView { return; } + if (!validateStopFrame(stopFrame, startFrame)) { + taskMessage.css('color', 'red'); + taskMessage.text('Stop frame must be greater than or equal to start frame'); + return; + } + if (files.length <= 0) { taskMessage.css('color', 'red'); taskMessage.text('No files specified for the task'); @@ -717,6 +761,15 @@ class DashboardView { if (customOverlapSize.prop('checked')) { description.overlap = overlapSize; } + if (customStartFrame.prop('checked')) { + description.start_frame = startFrame; + } + if (customStopFrame.prop('checked')) { + description.stop_frame = stopFrame; + } + if (customFrameFilter.prop('checked')) { + description.frame_filter = frameFilter; + } function cleanupTask(tid) { $.ajax({ diff --git a/cvat/apps/dashboard/templates/dashboard/dashboard.html b/cvat/apps/dashboard/templates/dashboard/dashboard.html index e75790d3..7722dc74 100644 --- a/cvat/apps/dashboard/templates/dashboard/dashboard.html +++ b/cvat/apps/dashboard/templates/dashboard/dashboard.html @@ -143,6 +143,33 @@ Example: @select=race:__undefined__,skip,asian,black,caucasian,other'/> + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/cvat/apps/engine/annotation.py b/cvat/apps/engine/annotation.py index 9c34548b..d0fb0211 100644 --- a/cvat/apps/engine/annotation.py +++ b/cvat/apps/engine/annotation.py @@ -1255,6 +1255,9 @@ class TaskAnnotation: ("flipped", str(db_task.flipped)), ("created", str(timezone.localtime(db_task.created_date))), ("updated", str(timezone.localtime(db_task.updated_date))), + ("start_frame", str(db_task.start_frame)), + ("stop_frame", str(db_task.stop_frame)), + ("frame_filter", db_task.frame_filter), ("labels", [ ("label", OrderedDict([ @@ -1416,7 +1419,7 @@ class TaskAnnotation: self._flip_shape(shape, im_w, im_h) dump_data = OrderedDict([ - ("frame", str(shape["frame"])), + ("frame", str(db_task.start_frame + shape["frame"] * db_task.get_frame_step())), ("outside", str(int(shape["outside"]))), ("occluded", str(int(shape["occluded"]))), ("keyframe", str(int(shape["keyframe"]))) diff --git a/cvat/apps/engine/migrations/0019_frame_selection.py b/cvat/apps/engine/migrations/0019_frame_selection.py new file mode 100644 index 00000000..d1b1d731 --- /dev/null +++ b/cvat/apps/engine/migrations/0019_frame_selection.py @@ -0,0 +1,40 @@ +# Generated by Django 2.1.7 on 2019-05-10 08:23 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('engine', '0018_jobcommit'), + ] + + operations = [ + migrations.RemoveField( + model_name='video', + name='start_frame', + ), + migrations.RemoveField( + model_name='video', + name='step', + ), + migrations.RemoveField( + model_name='video', + name='stop_frame', + ), + migrations.AddField( + model_name='task', + name='frame_filter', + field=models.CharField(default='', max_length=256), + ), + migrations.AddField( + model_name='task', + name='start_frame', + field=models.PositiveIntegerField(default=0), + ), + migrations.AddField( + model_name='task', + name='stop_frame', + field=models.PositiveIntegerField(default=0), + ), + ] diff --git a/cvat/apps/engine/models.py b/cvat/apps/engine/models.py index 736fd6b0..6a410d92 100644 --- a/cvat/apps/engine/models.py +++ b/cvat/apps/engine/models.py @@ -4,6 +4,7 @@ from enum import Enum +import re import shlex import os @@ -49,6 +50,9 @@ class Task(models.Model): z_order = models.BooleanField(default=False) flipped = models.BooleanField(default=False) image_quality = models.PositiveSmallIntegerField(default=50) + start_frame = models.PositiveIntegerField(default=0) + stop_frame = models.PositiveIntegerField(default=0) + frame_filter = models.CharField(max_length=256, default="") status = models.CharField(max_length=32, choices=StatusChoice.choices(), default=StatusChoice.ANNOTATION) @@ -64,6 +68,10 @@ class Task(models.Model): return path + def get_frame_step(self): + match = re.search("step\s*=\s*([1-9]\d*)", self.frame_filter) + return int(match.group(1)) if match else 1 + def get_upload_dirname(self): return os.path.join(self.get_task_dirname(), ".upload") @@ -128,9 +136,6 @@ class RemoteFile(models.Model): class Video(models.Model): task = models.OneToOneField(Task, on_delete=models.CASCADE) path = models.CharField(max_length=1024) - start_frame = models.PositiveIntegerField() - stop_frame = models.PositiveIntegerField() - step = models.PositiveIntegerField(default=1) width = models.PositiveIntegerField() height = models.PositiveIntegerField() diff --git a/cvat/apps/engine/serializers.py b/cvat/apps/engine/serializers.py index 7acc365b..85103449 100644 --- a/cvat/apps/engine/serializers.py +++ b/cvat/apps/engine/serializers.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: MIT import os +import re import shutil from rest_framework import serializers @@ -187,16 +188,25 @@ class TaskSerializer(WriteOnceMixin, serializers.ModelSerializer): fields = ('url', 'id', 'name', 'size', 'mode', 'owner', 'assignee', 'bug_tracker', 'created_date', 'updated_date', 'overlap', 'segment_size', 'z_order', 'flipped', 'status', 'labels', 'segments', - 'image_quality') + 'image_quality', 'start_frame', 'stop_frame', 'frame_filter') read_only_fields = ('size', 'mode', 'created_date', 'updated_date', 'status') write_once_fields = ('overlap', 'segment_size', 'image_quality') ordering = ['-id'] + def validate_frame_filter(self, value): + match = re.search("step\s*=\s*([1-9]\d*)", value) + if not match: + raise serializers.ValidationError("Invalid frame filter expression") + return value + # pylint: disable=no-self-use def create(self, validated_data): labels = validated_data.pop('label_set') db_task = models.Task.objects.create(size=0, **validated_data) + db_task.start_frame = validated_data.get('start_frame', 0) + db_task.stop_frame = validated_data.get('stop_frame', 0) + db_task.frame_filter = validated_data.get('frame_filter', '') for label in labels: attributes = label.pop('attributespec_set') db_label = models.Label.objects.create(task=db_task, **label) @@ -225,6 +235,9 @@ class TaskSerializer(WriteOnceMixin, serializers.ModelSerializer): instance.flipped = validated_data.get('flipped', instance.flipped) instance.image_quality = validated_data.get('image_quality', instance.image_quality) + instance.start_frame = validated_data.get('start_frame', instance.start_frame) + instance.stop_frame = validated_data.get('stop_frame', instance.stop_frame) + instance.frame_filter = validated_data.get('frame_filter', instance.frame_filter) labels = validated_data.get('label_set', []) for label in labels: attributes = label.pop('attributespec_set', []) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 5160c7d7..ca1992e4 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -7,6 +7,7 @@ import os import sys import rq import shutil +import subprocess import tempfile import numpy as np from PIL import Image @@ -48,15 +49,27 @@ def rq_handler(job, exc_type, exc_value, traceback): ############################# Internal implementation for server API class _FrameExtractor: - def __init__(self, source_path, compress_quality, flip_flag=False): + def __init__(self, source_path, compress_quality, step=1, start=0, stop=0, flip_flag=False): # translate inversed range 1:95 to 2:32 translated_quality = 96 - compress_quality translated_quality = round((((translated_quality - 1) * (31 - 2)) / (95 - 1)) + 2) + self.source = source_path self.output = tempfile.mkdtemp(prefix='cvat-', suffix='.data') target_path = os.path.join(self.output, '%d.jpg') output_opts = '-start_number 0 -b:v 10000k -vsync 0 -an -y -q:v ' + str(translated_quality) + filters = '' + if stop > 0: + filters = 'between(n,' + str(start) + ',' + str(stop) + ')' + elif start > 0: + filters = 'gte(n,' + str(start) + ')' + if step > 1: + filters += ('*' if filters else '') + 'not(mod(n-' + str(start) + ',' + str(step) + '))' + if filters: + filters = "select=\"'" + filters + "'\"" if flip_flag: - output_opts += ' -vf "transpose=2,transpose=2"' + filters += (',' if filters else '') + 'transpose=2,transpose=2' + if filters: + output_opts += ' -vf ' + filters ff = FFmpeg( inputs = {source_path: None}, outputs = {target_path: output_opts}) @@ -170,12 +183,13 @@ def _unpack_archive(archive, upload_dir): Archive(archive).extractall(upload_dir) os.remove(archive) -def _copy_video_to_task(video, db_task): +def _copy_video_to_task(video, db_task, step): job = rq.get_current_job() job.meta['status'] = 'Video is being extracted..' job.save_meta() - extractor = _FrameExtractor(video, db_task.image_quality) + extractor = _FrameExtractor(video, db_task.image_quality, + step, db_task.start_frame, db_task.stop_frame) for frame, image_orig_path in enumerate(extractor): image_dest_path = db_task.get_frame_path(frame) db_task.size += 1 @@ -183,10 +197,11 @@ def _copy_video_to_task(video, db_task): if not os.path.exists(dirname): os.makedirs(dirname) shutil.copyfile(image_orig_path, image_dest_path) + if db_task.stop_frame == 0: + db_task.stop_frame = db_task.start_frame + (db_task.size - 1) * step image = Image.open(db_task.get_frame_path(0)) models.Video.objects.create(task=db_task, path=video, - start_frame=0, stop_frame=db_task.size, step=1, width=image.width, height=image.height) image.close() @@ -351,7 +366,7 @@ def _create_thread(tid, data): if video: db_task.mode = "interpolation" video = os.path.join(upload_dir, video) - _copy_video_to_task(video, db_task) + _copy_video_to_task(video, db_task, db_task.get_frame_step()) else: db_task.mode = "annotation" _copy_images_to_task(upload_dir, db_task)