wehub-resource-sync
108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
# This script is intended to be used with an old release of CVAT, so it uses the internal
|
|
# API as it existed in that release. Because of this, we cannot use Pylint, since it will attempt
|
|
# to match the API calls against the current API and report spurious problems.
|
|
# pylint: disable=all
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import shutil
|
|
from unittest.mock import Mock, PropertyMock, patch
|
|
|
|
import django
|
|
from django.db import connection, transaction
|
|
from django.db.migrations.recorder import MigrationRecorder
|
|
from tqdm import tqdm
|
|
|
|
django.setup()
|
|
|
|
from cvat.apps.engine import models
|
|
from cvat.apps.engine.media_extractors import MEDIA_TYPES
|
|
from cvat.apps.engine.task import _create_static_chunks
|
|
|
|
EXPECTED_LAST_ENGINE_MIGRATION = "0095_fix_related_names"
|
|
|
|
|
|
def _ensure_last_engine_applied_migration_name():
|
|
recorder = MigrationRecorder(connection)
|
|
app_name = "engine"
|
|
applied_migrations_names = list(
|
|
recorder.Migration.objects.filter(app=app_name).values_list("name", flat=True)
|
|
)
|
|
assert applied_migrations_names, f"No migrations applied for app '{app_name}'"
|
|
|
|
highest_by_number = max(applied_migrations_names, key=lambda name: int(name.split("_")[0]))
|
|
|
|
assert highest_by_number == EXPECTED_LAST_ENGINE_MIGRATION, (
|
|
f"Last applied migration for app '{app_name}' is '{highest_by_number}', "
|
|
f"expected '{EXPECTED_LAST_ENGINE_MIGRATION}'. "
|
|
f"Manually verify that the script still works correctly "
|
|
f"and update EXPECTED_ENGINE_LAST_ENGINE_MIGRATION var in the script."
|
|
)
|
|
|
|
|
|
def _build_extractor(data: models.Data):
|
|
details = {
|
|
"source_path": [os.path.join(data.get_upload_dirname(), data.video.path)],
|
|
"step": data.get_frame_step(),
|
|
"start": data.start_frame,
|
|
"stop": data.stop_frame,
|
|
}
|
|
return MEDIA_TYPES["video"]["extractor"](**details)
|
|
|
|
|
|
def _cleanup_static_cache(data: models.Data):
|
|
for folder in (data.get_compressed_cache_dirname(), data.get_original_cache_dirname()):
|
|
if os.path.exists(folder):
|
|
shutil.rmtree(folder)
|
|
os.makedirs(folder)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Switch a video task to static cache storage and generate static chunks."
|
|
)
|
|
parser.add_argument("task_id", type=int, help="Task ID (integer)")
|
|
args = parser.parse_args()
|
|
task_id = args.task_id
|
|
|
|
_ensure_last_engine_applied_migration_name()
|
|
|
|
with transaction.atomic():
|
|
try:
|
|
task: models.Task = models.Task.objects.select_for_update().get(pk=task_id)
|
|
except models.Task.DoesNotExist as ex:
|
|
raise ValueError(f"Task #{task_id} not found.") from ex
|
|
|
|
assert (
|
|
task.mode == "interpolation"
|
|
), f"Task #{task_id} is not a video task (mode={task.mode})."
|
|
|
|
data: models.Data = task.data
|
|
|
|
extractor = _build_extractor(data)
|
|
|
|
data.storage_method = models.StorageMethodChoice.FILE_SYSTEM
|
|
_cleanup_static_cache(data)
|
|
|
|
with tqdm(total=1.0, bar_format="{l_bar}{bar}| {n:.2f}/{total:.2f}") as pbar:
|
|
|
|
class _PropertyMock(PropertyMock):
|
|
def __set__(self, instance, value):
|
|
pbar.n = value
|
|
pbar.refresh()
|
|
|
|
with patch("cvat.apps.engine.task.ImportRQMeta", return_value=Mock()) as mock:
|
|
type(mock.for_job.return_value).task_progress = _PropertyMock()
|
|
_create_static_chunks(
|
|
task, media_extractor=extractor, upload_dir=data.get_upload_dirname()
|
|
)
|
|
data.save(update_fields=["storage_method"])
|
|
|
|
print(f"Task #{task_id}: switched to static cache.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|