forked from prehistoric-systems/comixify
b5dd5cffc0
* Add keyframe model * Add segmentation utils * Add keyframes extraction pipeline * Add keyframe tests * Update dockerfile to include caffe * Add summe pretrained model * Add video for testing * Update keyframe pipeline, tests * Update settings to use in memory db for tests * Set keyframe number to 10, fix bugs * Fix keyframe order * Make requested changes * Fix Dockerfile * Make requested changes * Make requested changes * Add blank lines * Change dockerfile base cuda image to devel version * Add modified Cuda.cmake for Dockerfile * Add pyyaml dependecy to dockerfile * Update dockerfile * Update dockerfile * Fix markdown version error * Fix markdown version error * Change caffe installation to make * Update dockerfile * Update dockerfile * Fix boost imoprt * Fix boost not found bug * Add feature normalisation * Fix dateutil, fix caffe root slash * Fix slash bug * Add batching to feature extraction * Add model caching to keyframes extraction * Fix output images to be in proper range * Add time logging * Change feature batch to 128 * Change dockerfile * Fix dockerfile * Change feature batch to 10 * Add set mode gpu * Change feature batch to 64 * Change feature batch to 32 * Add I-frame frame sampling * Cleanup * Delete Cuda.cmake * Remove comments from Makefile.config * Cleanup * Fix color scheme switching * Remove cudnn.hpp, change caffe to 1.0 * Remove cudnn.hpp copy in dockerfile * Remove redundant run's in dockerfile * Change pretrained model
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
from django.test import TestCase
|
|
|
|
import numpy as np
|
|
from keyframes.keyframes import KeyFramesExtractor
|
|
from api.models import Video
|
|
from django.core.files import File
|
|
import shutil
|
|
from utils import jj
|
|
from django.conf import settings
|
|
from keyframes.utils import batch
|
|
|
|
VIDEO_PATH = "tmp/f1_short.mp4"
|
|
VIDEO_N_FRAMES = 47
|
|
|
|
|
|
class KeyframesTestCase(TestCase):
|
|
|
|
def setUp(self):
|
|
f = open(VIDEO_PATH, 'rb')
|
|
self.video = Video.objects.create(file=File(f))
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(jj(f"{settings.TMP_DIR}", f"{self.all_frames_tmp_dir}"))
|
|
|
|
def test_keyframes(self):
|
|
"""Keyframes are extracted corectly"""
|
|
|
|
frames_paths, all_frames_tmp_dir = KeyFramesExtractor._get_all_frames(self.video)
|
|
self.assertIsInstance(frames_paths[0], str)
|
|
self.assertEqual(len(frames_paths), VIDEO_N_FRAMES)
|
|
self.all_frames_tmp_dir = all_frames_tmp_dir
|
|
|
|
frames = KeyFramesExtractor._get_frames(frames_paths)
|
|
self.assertEqual(len(frames), VIDEO_N_FRAMES)
|
|
self.assertIsInstance(frames[0], np.ndarray)
|
|
|
|
features = KeyFramesExtractor._get_features(frames, False)
|
|
self.assertIsInstance(features, np.ndarray)
|
|
self.assertEqual(features.shape, (VIDEO_N_FRAMES, 1024))
|
|
|
|
change_points, frames_per_segment = KeyFramesExtractor._get_segments(features)
|
|
self.assertIsInstance(change_points, list)
|
|
self.assertIsInstance(frames_per_segment, list)
|
|
|
|
for cp in frames_per_segment:
|
|
with self.subTest(cp=cp):
|
|
self.assertIsInstance(cp, int)
|
|
|
|
probs = KeyFramesExtractor._get_probs(features, False)
|
|
self.assertIsInstance(probs, np.ndarray)
|
|
self.assertEqual(probs.shape, (VIDEO_N_FRAMES, ))
|
|
|
|
chosen_frames = KeyFramesExtractor._get_chosen_frames(frames, probs, change_points, frames_per_segment)
|
|
self.assertIsInstance(chosen_frames, list)
|
|
self.assertTrue(len(chosen_frames) == 10)
|
|
|
|
|
|
class UtilsTestCase(TestCase):
|
|
def test_batch(self):
|
|
"""Barch is working"""
|
|
arr = [1, 1, 2, 2, 3, 3, 4]
|
|
batched_arr = batch(arr, 2)
|
|
self.assertEqual(list(batched_arr), [(2, [1, 1]), (2, [2, 2]), (2, [3, 3]), (1, [4])])
|
|
|
|
def test_empty_batch(self):
|
|
"""Barch is working"""
|
|
arr = []
|
|
batched_arr = batch(arr, 2)
|
|
self.assertEqual(list(batched_arr), [])
|