Skip to content

Commit

Permalink
Support frame selection when create from video
Browse files Browse the repository at this point in the history
  • Loading branch information
zliang7 committed May 31, 2019
1 parent 80e8a7b commit b8f8257
Show file tree
Hide file tree
Showing 8 changed files with 168 additions and 11 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
53 changes: 53 additions & 0 deletions cvat/apps/dashboard/static/dashboard/js/dashboard.js
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand Down Expand Up @@ -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');
Expand All @@ -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', () => {
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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');
Expand Down Expand Up @@ -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');
Expand Down Expand Up @@ -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({
Expand Down
27 changes: 27 additions & 0 deletions cvat/apps/dashboard/templates/dashboard/dashboard.html
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,33 @@
<input type="checkbox" id="dashboardCustomQuality" title="Custom image quality"/>
</td>
</tr>
<tr>
<td>
<label class="regular h2"> Start Frame </label>
</td>
<td>
<input type="number" id="dashboardStartFrame" class="regular" style="width: 4.5em;" min="0" value=0 disabled=true/>
<input type="checkbox" id="dashboardCustomStart" title="Custom start frame"/>
</td>
</tr>
<tr>
<td>
<label class="regular h2"> Stop Frame </label>
</td>
<td>
<input type="number" id="dashboardStopFrame" class="regular" style="width: 4.5em;" min="0" value=0 disabled=true/>
<input type="checkbox" id="dashboardCustomStop" title="Custom stop frame"/>
</td>
</tr>
<tr>
<td>
<label class="regular h2"> Frame Filter </label>
</td>
<td>
<input type="text" id="dashboardFrameFilter" class="regular" style="width: 4.5em;" title="Currently only support 'step=K' filter expression." disabled=true/>
<input type="checkbox" id="dashboardCustomFilter" title="Custom frame filter"/>
</td>
</tr>
</table>


Expand Down
5 changes: 4 additions & 1 deletion cvat/apps/engine/annotation.py
Original file line number Diff line number Diff line change
Expand Up @@ -1255,6 +1255,9 @@ def dump(self, file_path, scheme, host, query_params):
("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([
Expand Down Expand Up @@ -1414,7 +1417,7 @@ def dump(self, file_path, scheme, host, query_params):
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"])))
Expand Down
40 changes: 40 additions & 0 deletions cvat/apps/engine/migrations/0019_frame_selection.py
Original file line number Diff line number Diff line change
@@ -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),
),
]
11 changes: 8 additions & 3 deletions cvat/apps/engine/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from enum import Enum

import re
import shlex
import os

Expand Down Expand Up @@ -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)

Expand All @@ -64,6 +68,10 @@ def get_frame_path(self, frame):

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")

Expand Down Expand Up @@ -128,9 +136,6 @@ class Meta:
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()

Expand Down
15 changes: 14 additions & 1 deletion cvat/apps/engine/serializers.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
# SPDX-License-Identifier: MIT

import os
import re
import shutil

from rest_framework import serializers
Expand Down Expand Up @@ -187,16 +188,25 @@ class Meta:
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)
Expand Down Expand Up @@ -225,6 +235,9 @@ def update(self, instance, validated_data):
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', [])
Expand Down
27 changes: 21 additions & 6 deletions cvat/apps/engine/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import sys
import rq
import shutil
import subprocess
import tempfile
import numpy as np
from PIL import Image
Expand Down Expand Up @@ -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})
Expand Down Expand Up @@ -170,23 +183,25 @@ 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
dirname = os.path.dirname(image_dest_path)
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()

Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit b8f8257

Please sign in to comment.