From 2e2ca583906c61e0d53a4acf78710eb3cbe946c3 Mon Sep 17 00:00:00 2001 From: Bryon Lewis Date: Wed, 12 Aug 2026 10:20:52 -0400 Subject: [PATCH] refactor and split tasks --- server/dive_tasks/__init__.py | 5 + server/dive_tasks/convert_images.py | 372 +++++ server/dive_tasks/convert_video.py | 277 ++++ server/dive_tasks/run_pipeline.py | 574 +++++++ server/dive_tasks/run_training.py | 207 +++ server/dive_tasks/tasks.py | 1656 +------------------- server/dive_tasks/upgrade_pipelines.py | 140 ++ server/dive_tasks/viame_config.py | 100 ++ server/tests/test_google_drive_download.py | 4 +- server/tests/test_inject_metadata_file.py | 2 +- server/tests/test_remote_ffprobe.py | 32 +- 11 files changed, 1742 insertions(+), 1627 deletions(-) create mode 100644 server/dive_tasks/convert_images.py create mode 100644 server/dive_tasks/convert_video.py create mode 100644 server/dive_tasks/run_pipeline.py create mode 100644 server/dive_tasks/run_training.py create mode 100644 server/dive_tasks/upgrade_pipelines.py create mode 100644 server/dive_tasks/viame_config.py diff --git a/server/dive_tasks/__init__.py b/server/dive_tasks/__init__.py index 48e53f38f..6a13743dc 100644 --- a/server/dive_tasks/__init__.py +++ b/server/dive_tasks/__init__.py @@ -15,6 +15,11 @@ def task_imports(self): # worker_girder_events first: bind Girder handlers before any task module loads. return [ 'dive_tasks.worker_girder_events', + 'dive_tasks.upgrade_pipelines', + 'dive_tasks.run_pipeline', + 'dive_tasks.run_training', + 'dive_tasks.convert_video', + 'dive_tasks.convert_images', 'dive_tasks.tasks', 'dive_tasks.local_tasks', ] diff --git a/server/dive_tasks/convert_images.py b/server/dive_tasks/convert_images.py new file mode 100644 index 000000000..202dbcfc2 --- /dev/null +++ b/server/dive_tasks/convert_images.py @@ -0,0 +1,372 @@ +from contextlib import suppress +import os +from pathlib import Path +import shlex +import tempfile +import zipfile + +from girder_client import GirderClient, HttpError +from girder_worker.app import app +from girder_worker.task import Task +from girder_worker.utils import JobManager, JobStatus + +from dive_tasks import utils +from dive_tasks.convert_video import resolve_annotation_fps +from dive_tasks.manager import patch_manager +from dive_tasks.viame_config import Config +from dive_utils import asbool, calibration_format, constants +from dive_utils.types import GirderModel + + +@app.task(bind=True, acks_late=True, ignore_result=True) +def convert_calibration(self: Task, itemId: str): + """ + Convert a calibrationFile item to a JSON camera-rig in a separate Girder item + marked jsonCalibrationFile for display. + """ + conf = Config() + conf.require_viame_install() + context: dict = {} + gc: GirderClient = self.girder_client + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + convert_tool = conf.viame_install_path / 'configs' / 'convert_cam_format.py' + + with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): + _working_directory_path = Path(_working_directory) + item: GirderModel = gc.getItem(itemId) + folder_id = str(item.get('folderId')) + folder = gc.getFolder(folder_id) + multi_cam = (folder.get('meta') or {}).get(constants.MultiCamMarker) or {} + if str(multi_cam.get(constants.CalibrationItemIdMarker)) != str(itemId): + manager.write('Calibration source was replaced; skipping stale conversion job.\n') + return + + files = list(gc.listFile(itemId)) + source_file = next( + (f for f in files if constants.stereoCalibrationRegex.search(f['name'])), + None, + ) + if source_file is None: + manager.write('No convertible calibration file found in item; skipping.\n') + return + + manager.updateStatus(JobStatus.FETCHING_INPUT) + gc.downloadItem(itemId, _working_directory_path, name=item.get('name')) + input_path = _working_directory_path / source_file['name'] + file_bytes = input_path.read_bytes() + if source_file['name'].lower().endswith('.json'): + if calibration_format.calibration_upload_is_final_json(source_file['name'], file_bytes): + manager.write('Source calibration is already JSON; skipping conversion.\n') + return + input_path = calibration_format.prepare_conversion_input_path( + input_path, + file_bytes, + _working_directory_path, + ) + output_path = input_path.with_suffix('.json') + + manager.updateStatus(JobStatus.RUNNING) + command = [ + f". {shlex.quote(str(conf.viame_setup_script))} &&", + "python", + shlex.quote(str(convert_tool)), + shlex.quote(str(input_path)), + shlex.quote(str(output_path)), + ] + popen_kwargs = { + 'args': " ".join(command), + 'shell': True, + 'executable': '/bin/bash', + 'cwd': str(_working_directory_path), + 'env': conf.gpu_process_env, + } + try: + utils.stream_subprocess(self, context, manager, popen_kwargs) + except Exception as exc: + error_msg = str(exc) or 'Calibration conversion failed' + updated_multi_cam = dict(multi_cam) + updated_multi_cam[constants.CalibrationConversionErrorMarker] = error_msg + gc.addMetadataToFolder( + folder_id, + {constants.MultiCamMarker: updated_multi_cam}, + ) + raise + + if not output_path.exists() or not output_path.stat().st_size: + error_msg = 'Calibration conversion produced no JSON output' + updated_multi_cam = dict(multi_cam) + updated_multi_cam[constants.CalibrationConversionErrorMarker] = error_msg + gc.addMetadataToFolder( + folder_id, + {constants.MultiCamMarker: updated_multi_cam}, + ) + raise RuntimeError(error_msg) + + folder = gc.getFolder(folder_id) + multi_cam = (folder.get('meta') or {}).get(constants.MultiCamMarker) or {} + if str(multi_cam.get(constants.CalibrationItemIdMarker)) != str(itemId): + manager.write('Calibration source was replaced during conversion; discarding output.\n') + return + + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + json_name = output_path.name + gc.upload(str(output_path), folder_id) + json_items = sorted( + gc.listItem(folder_id, name=json_name), + key=lambda existing: existing.get('created', ''), + ) + if not json_items: + raise RuntimeError('Failed to create calibration JSON item') + json_item = json_items[-1] + json_item_id = str(json_item['_id']) + gc.addMetadataToItem( + json_item_id, + { + constants.JsonCalibrationFileMarker: 'true', + constants.CalibrationFileMarker: False, + }, + ) + + folder = gc.getFolder(folder_id) + multi_cam = (folder.get('meta') or {}).get(constants.MultiCamMarker) or {} + if str(multi_cam.get(constants.CalibrationItemIdMarker)) != str(itemId): + manager.write( + 'Calibration source was replaced before linking JSON; discarding output.\n' + ) + gc.delete(f'item/{json_item_id}') + return + + updated_multi_cam = dict(multi_cam) + updated_multi_cam.setdefault(constants.CalibrationItemIdMarker, str(itemId)) + updated_multi_cam[constants.JsonCalibrationItemIdMarker] = json_item_id + updated_multi_cam.setdefault( + constants.CalibrationOriginalNameMarker, + source_file['name'], + ) + updated_multi_cam.pop(constants.CalibrationConversionErrorMarker, None) + gc.addMetadataToFolder( + folder_id, + {constants.MultiCamMarker: updated_multi_cam}, + ) + + for existing_item in gc.listItem(folder_id): + existing_id = str(existing_item['_id']) + if existing_id in {json_item_id, str(itemId)}: + continue + existing_meta = existing_item.get('meta') or {} + if asbool(existing_meta.get(constants.JsonCalibrationFileMarker)): + gc.delete(f"item/{existing_id}") + + +@app.task(bind=True, acks_late=True) +def convert_images(self: Task, folderId, user_id: str, user_login: str): + """ + Ensures that all images in a folder are in a web friendly format (png or jpeg). + + If conversions succeeds for an image, it will replace the image with an image + of the same name, but in a web friendly extension. + + Returns the number of images successfully converted. + """ + context: dict = {} + gc: GirderClient = self.girder_client + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + items_to_convert = [ + item + for item in gc.listItem(folderId) + if ( + constants.imageRegex.search(item["name"]) + and not constants.safeImageRegex.search(item["name"]) + ) + ] + + with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): + working_directory_path = Path(_working_directory) + images_path = utils.make_directory(working_directory_path / 'images') + + for item in items_to_convert: + # Assumes 1 file per item + gc.downloadItem(item["_id"], images_path, item["name"]) + + item_path = images_path / item["name"] + new_item_path = images_path / ".".join([*item["name"].split(".")[:-1], "png"]) + command = ["ffmpeg", "-i", str(item_path), str(new_item_path)] + utils.stream_subprocess(self, context, manager, {'args': command}) + gc.uploadFileToFolder(folderId, new_item_path) + gc.delete(f"item/{str(item['_id'])}") + + gc.addMetadataToFolder( + str(folderId), + { + "annotate": True, # mark the parent folder as able to annotate. + constants.FPSMarker: resolve_annotation_fps(gc, folderId), + }, + ) + + +@app.task(bind=True, acks_late=True) +def convert_large_images(self: Task, folderId, user_id: str, user_login: str): + """ + Converts all images in the folder to large images + + This is typically done if the images are >8k W or L resolution + + Returns the number of images successfully converted. + """ + context: dict = {} + gc: GirderClient = self.girder_client + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + items_to_convert = [ + item for item in gc.listItem(folderId) if (constants.safeImageRegex.search(item["name"])) + ] + for item in items_to_convert: + # Assumes 1 file per item + try: + # Does it already have tiles? + gc.get(f'item/{item["_id"]}/tiles') + manager.write(f'Skipping {item["name"]}, already a large image\n') + continue + except HttpError as e: + # Safely parse JSON if possible + message = "" + try: + message = e.response.json().get("message", "") + except Exception: + pass # non-JSON response, leave message empty + # This is the Girder message when no large image exists + if e.status == 400 and message == "No large image file in this item.": + manager.write(f'Converting {item["name"]} to large image\n') + gc.post(f'item/{item["_id"]}/tiles') + else: + # Re-raise unexpected errors to fail the job + raise + gc.addMetadataToFolder( + str(folderId), + {"type": constants.LargeImageType}, # mark the parent folder as able to annotate. + ) + + +@app.task(bind=True, acks_late=True, ignore_result=True) +def extract_zip(self: Task, folderId: str, itemId: str, user_id: str, user_login: str): + """ + Discovery logic: + * Find all folders that have at least one child file (potential datasets) + * Exclude folders which are sub-folders of previously discovered folders + because datasets cannot be nested in other datasets + """ + context: dict = {} + gc: GirderClient = self.girder_client + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): + _working_directory_path = Path(_working_directory) + item: GirderModel = gc.getItem(itemId) + file_name = str(_working_directory_path / item['name']) + manager.write(f'Fetching input from {itemId} to {file_name}...\n') + gc.downloadItem(itemId, _working_directory, item["name"]) + discovered_folders = {} + with zipfile.ZipFile(file_name, 'r') as zipObj: + listOfFileNames = zipObj.namelist() + sum_file_size = sum([data.file_size for data in zipObj.filelist]) + sum_compress_size = sum([data.compress_size for data in zipObj.filelist]) + ratio = sum_file_size / sum_compress_size + if ratio > 600: + manager.write(f"Compression ratio is exceedingly high at {ratio}\n\ + Please contact an admin at viame-web@kitware.com if this is a valid zip file") + raise Exception("High Compression Ratio for Zip File") + + multicam_export_roots = { + os.path.dirname(fileName) + for fileName in listOfFileNames + if os.path.basename(fileName) == constants.MultiCamJsonFileName + and not fileName.endswith(os.path.sep) + } + + for fileName in listOfFileNames: + folderName = os.path.dirname(fileName) + parentName = os.path.dirname(folderName) + if parentName in discovered_folders and folderName != '': + discovered_folders[folderName] = 'ignored' + # Nested single-camera exports stay skipped; multicam camera trees must extract. + if not utils.is_path_under_multicam_export(folderName, multicam_export_roots): + continue + if fileName.endswith(os.path.sep): + continue + if folderName not in discovered_folders: + discovered_folders[folderName] = 'unstructured' + if constants.metaRegex.search(os.path.basename(fileName)): + if folderName in multicam_export_roots: + discovered_folders[folderName] = 'multicam' + else: + discovered_folders[folderName] = 'dataset' + if fileName.endswith('.zip'): + raise Exception("Nested Zip Files are invalid") + manager.write(f"Extracting: {fileName}\n") + zipObj.extract(fileName, f'{_working_directory}') + + # remove the zip file so it isn't uploaded back to the folder + os.remove(file_name) + # Create source folder and move zip file there + created_folder = gc.createFolder( + folderId, + constants.SourceFolderName, + reuseExisting=True, + ) + gc.sendRestRequest( + "PUT", + f"/item/{str(item['_id'])}?folderId={str(created_folder['_id'])}", + ) + # Only make subfolders if more than 1 discovered folder exists + make_subfolders = ( + len(discovered_folders) - list(discovered_folders.values()).count('ignored') + ) > 1 + for folderName, folderType in discovered_folders.items(): + subFolderName = folderName if make_subfolders else '' + if folderType == 'unstructured': + utils.upload_zipped_flat_media_files( + gc, + manager, + folderId, + _working_directory_path / folderName, + subFolderName, + ) + elif folderType == 'multicam': + utils.upload_exported_multicam_zipped_dataset( + gc, + manager, + folderId, + _working_directory_path / folderName, + subFolderName, + ) + elif folderType == 'dataset': + utils.upload_exported_zipped_dataset( + gc, + manager, + folderId, + _working_directory_path / folderName, + subFolderName, + ) + else: + manager.write(f'Ignoring {folderName}\n') + + if make_subfolders: + gc.sendRestRequest( + "DELETE", + f"folder/{folderId}/metadata", + json=[constants.TypeMarker, constants.FPSMarker, constants.DatasetMarker], + ) diff --git a/server/dive_tasks/convert_video.py b/server/dive_tasks/convert_video.py new file mode 100644 index 000000000..1cd7937ac --- /dev/null +++ b/server/dive_tasks/convert_video.py @@ -0,0 +1,277 @@ +from contextlib import suppress +from pathlib import Path +import tempfile +from typing import Optional + +from girder_client import GirderClient +from girder_worker.app import app +from girder_worker.task import Task +from girder_worker.utils import JobManager, JobStatus + +from dive_tasks import utils +from dive_tasks.frame_alignment import check_and_fix_frame_alignment, is_frame_misaligned +from dive_tasks.manager import patch_manager +from dive_utils import constants, fromMeta +from dive_utils.types import GirderModel + + +def resolve_annotation_fps( + gc: GirderClient, + folder_id: str, + *, + native_fps: Optional[float] = None, + default_fps: float = 1.0, +) -> float: + """Pick annotation FPS from current folder meta vs media FPS. + + Re-reads folder ``fps`` so a concurrent CSV import (assetstore postprocess) + is not overwritten by a stale ``-1`` snapshot from job start. + + For video, pass ``native_fps`` from ffprobe. For image sequences, omit it so + ``-1`` falls back to ``default_fps`` (1) and any CSV-set value is kept. + """ + requested_fps = fromMeta(gc.getFolder(folder_id), constants.FPSMarker) + return utils.choose_annotation_fps( + requested_fps, native_fps=native_fps, default_fps=default_fps + ) + + +def _download_video_item( + gc: GirderClient, + manager: JobManager, + item_id: str, + item_name: str, + dest_dir: Path, +) -> str: + """Download a Girder video item to *dest_dir*; return the local file path.""" + file_name = str(dest_dir / item_name) + manager.updateStatus(JobStatus.FETCHING_INPUT) + manager.write(f'Fetching input from {item_id} to {file_name}...\n') + gc.downloadItem(item_id, dest_dir, name=item_name) + return file_name + + +@app.task(bind=True, acks_late=True, ignore_result=True) +def convert_video( + self: Task, folderId: str, itemId: str, user_id: str, user_login: str, skip_transcoding=False +): + context: dict = {} + gc: GirderClient = self.girder_client + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): + _working_directory_path = Path(_working_directory) + item: GirderModel = gc.getItem(itemId) + item_name = item['name'] + output_file_path = (_working_directory_path / item_name).with_suffix('.transcoded.mp4') + + # When skip_transcoding is requested, probe via authenticated HTTP Range + # requests first so web-ready S3/filesystem videos never need a full download. + # Fall back to downloading the whole object if remote probe/alignment fails. + jsoninfo = None + file_name: Optional[str] = None + auth_headers: Optional[str] = None + remote_url: Optional[str] = None + + if skip_transcoding: + try: + remote_url = utils.item_primary_file_download_url(gc, itemId) + except Exception as exc: + manager.write( + f'Could not resolve file download URL ({exc}); ' + 'falling back to full download\n' + ) + remote_url = None + if remote_url is not None: + auth_headers = utils.girder_auth_headers(gc.token) + manager.updateStatus(JobStatus.RUNNING) + manager.write(f'Probing video via HTTP Range requests: {remote_url}\n') + try: + jsoninfo = utils.ffprobe_format_and_streams( + self, context, manager, remote_url, headers=auth_headers + ) + except utils.CanceledError: + raise + except Exception as exc: + manager.write(f'Remote ffprobe failed ({exc}); falling back to full download\n') + jsoninfo = None + auth_headers = None + remote_url = None + + if jsoninfo is None: + file_name = _download_video_item( + gc, manager, itemId, item_name, _working_directory_path + ) + manager.updateStatus(JobStatus.RUNNING) + jsoninfo = utils.ffprobe_format_and_streams(self, context, manager, file_name) + + videostream = list(filter(lambda x: x["codec_type"] == "video", jsoninfo["streams"])) + if len(videostream) != 1: + print('Expected 1 video stream, found {}'.format(len(videostream))) + print('Using first Video Stream found') + + format_info = jsoninfo.get('format') or {} + format_name = format_info.get('format_name') or '' + + # Extract framerate (avg_frame_rate, else r_frame_rate for e.g. MPEG-TS) + originalFpsString, originalFps = utils.fps_from_ffprobe_stream(videostream[0]) + + source_misaligned = False + if skip_transcoding: + alignment_source: Optional[str] = file_name or remote_url + alignment_headers = auth_headers if file_name is None else None + try: + source_misaligned = is_frame_misaligned( + self, + alignment_source, + context, + manager, + headers=alignment_headers, + ) + except utils.CanceledError: + raise + except Exception as exc: + if file_name is None: + manager.write( + f'Remote frame-alignment check failed ({exc}); ' + 'falling back to full download\n' + ) + file_name = _download_video_item( + gc, manager, itemId, item_name, _working_directory_path + ) + manager.updateStatus(JobStatus.RUNNING) + # Re-probe locally so metadata matches the file we will encode. + jsoninfo = utils.ffprobe_format_and_streams(self, context, manager, file_name) + videostream = list( + filter(lambda x: x["codec_type"] == "video", jsoninfo["streams"]) + ) + format_info = jsoninfo.get('format') or {} + format_name = format_info.get('format_name') or '' + originalFpsString, originalFps = utils.fps_from_ffprobe_stream(videostream[0]) + source_misaligned = is_frame_misaligned(self, Path(file_name), context, manager) + else: + raise + + # Skip remux/transcode only for browser-safe sources, matching desktop checks. + can_skip_transcode = utils.can_skip_video_transcoding( + skip_transcoding=skip_transcoding, + codec_name=videostream[0]['codec_name'], + sample_aspect_ratio=videostream[0].get('sample_aspect_ratio'), + format_name=format_name, + source_misaligned=source_misaligned, + ) + + # lets determine if we don't need to transcode this file + if can_skip_transcode: + # Now we can update the meta data and push the values + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + if file_name is None: + manager.write('Skip transcode: no full download required\n') + newAnnotationFps = resolve_annotation_fps(gc, folderId, native_fps=originalFps) + gc.addMetadataToItem( + itemId, + { + "source_video": False, # even though it is, this for requesting + "transcoder": "ffmpeg", + constants.OriginalFPSMarker: originalFps, + constants.OriginalFPSStringMarker: originalFpsString, + "codec": "h264", + }, + ) + gc.addMetadataToFolder( + folderId, + { + constants.DatasetMarker: True, # mark the parent folder as able to annotate. + constants.OriginalFPSMarker: originalFps, + constants.OriginalFPSStringMarker: originalFpsString, + constants.FPSMarker: newAnnotationFps, + "ffprobe_info": videostream[0], + }, + ) + return + elif skip_transcoding: + print('Transcoding cannot be skipped:') + print(f'Codec Name: {videostream[0]["codec_name"]}') + print(f'format_name: {format_name}') + if videostream[0]['codec_name'] != 'h264': + print('Codec is not h264; file will be transcoded') + elif videostream[0].get('sample_aspect_ratio') != '1:1': + print( + 'Sample aspect ratio is not 1:1; file will be transcoded ' + '(desktop-parity rule)' + ) + elif not utils.container_allows_skip_transcoding(format_name): + print('Container is not web-safe (e.g. mpegts); file will be transcoded') + elif source_misaligned: + print('Frame timestamps are misaligned; file will be transcoded') + + if file_name is None: + file_name = _download_video_item( + gc, manager, itemId, item_name, _working_directory_path + ) + manager.updateStatus(JobStatus.RUNNING) + + command = [ + "ffmpeg", + "-i", + file_name, + "-c:v", + "libx264", + "-preset", + "slow", + # https://github.com/Kitware/dive/issues/855 + "-crf", + "22", + # https://askubuntu.com/questions/1315697/could-not-find-tag-for-codec-pcm-s16le-in-stream-1-codec-not-currently-support + "-c:a", + "aac", + # see native/ code for a discussion of this option + "-vf", + "scale=ceil(iw*sar/2)*2:ceil(ih/2)*2,setsar=1", + str(output_file_path), + ] + utils.stream_subprocess(self, context, manager, {'args': command}) + # Check to see if frame alignment remains the same + aligned_file = check_and_fix_frame_alignment(self, output_file_path, context, manager) + misaligned_flag = False + if aligned_file != output_file_path: + misaligned_flag = True + + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + newAnnotationFps = resolve_annotation_fps(gc, folderId, native_fps=originalFps) + new_file = gc.uploadFileToFolder(folderId, aligned_file) + gc.addMetadataToItem( + new_file['itemId'], + { + "source_video": False, + "transcoder": "ffmpeg", + constants.OriginalFPSMarker: originalFps, + constants.OriginalFPSStringMarker: originalFpsString, + "codec": "h264", + }, + ) + source_metadata = { + "source_video": True, + constants.OriginalFPSMarker: originalFps, + constants.OriginalFPSStringMarker: originalFpsString, + "codec": videostream[0]["codec_name"], + } + if misaligned_flag: + source_metadata[constants.MISALGINED_MARKER] = True + gc.addMetadataToItem( + itemId, + source_metadata, + ) + gc.addMetadataToFolder( + folderId, + { + constants.DatasetMarker: True, # mark the parent folder as able to annotate. + constants.OriginalFPSMarker: originalFps, + constants.OriginalFPSStringMarker: originalFpsString, + constants.FPSMarker: newAnnotationFps, + "ffprobe_info": videostream[0], + }, + ) diff --git a/server/dive_tasks/run_pipeline.py b/server/dive_tasks/run_pipeline.py new file mode 100644 index 000000000..ab3a4aab2 --- /dev/null +++ b/server/dive_tasks/run_pipeline.py @@ -0,0 +1,574 @@ +from contextlib import suppress +import os +from pathlib import Path +import shlex +import shutil +import tempfile +from typing import Dict, List, Optional, Tuple + +from girder_client import GirderClient +from girder_worker.app import app +from girder_worker.task import Task +from girder_worker.utils import JobManager, JobStatus + +from dive_tasks import utils +from dive_tasks.manager import patch_manager +from dive_tasks.multicam_pipeline import ( + append_metadata_file_kwiver_settings, + append_stereo_calibration_kwiver_settings, + build_multicam_kwiver_settings, + find_downloaded_calibration_file, + is_stereo_measurement_pipeline, +) +from dive_tasks.pipeline_creates_dataset import ( + append_new_dataset_media_writers, + is_transcode_pipeline, + pipeline_creates_new_dataset, + pipeline_renumbers_frames, +) +from dive_tasks.viame_config import Config +from dive_utils import constants, fromMeta +from dive_utils.types import GirderModel, MulticamCameraJob, MulticamPipelineJob, PipelineJob + + +def filter_csv_by_frame_range(csv_path: str, frame_range: Tuple[int, int]) -> str: + """Filter VIAME CSV to only include detections within frame range. + + Args: + csv_path: Path to the input CSV file + frame_range: Tuple of (start_frame, end_frame) inclusive + + Returns: + Path to the filtered CSV file + """ + start_frame, end_frame = frame_range + filtered_path = csv_path.replace('.csv', '_filtered.csv') + + with open(csv_path, 'r') as infile, open(filtered_path, 'w') as outfile: + for line in infile: + if line.startswith('#'): + outfile.write(line) + continue + parts = line.split(',') + if len(parts) >= 3: + try: + frame = int(parts[2]) # Frame number is column 3 (0-indexed as column 2) + if start_frame <= frame <= end_frame: + outfile.write(line) + except ValueError: + # If frame number can't be parsed, include the line + outfile.write(line) + return filtered_path + + +def filter_image_list_by_frame_range( + image_list: List[str], frame_range: Tuple[int, int] +) -> List[str]: + """Filter an image list to only include images within frame range. + + Args: + image_list: List of image file paths + frame_range: Tuple of (start_frame, end_frame) inclusive (0-indexed) + + Returns: + Filtered list of image file paths + """ + start_frame, end_frame = frame_range + # Ensure we don't go out of bounds + start_frame = max(0, start_frame) + end_frame = min(end_frame, len(image_list) - 1) + return image_list[start_frame : end_frame + 1] + + +def _resolve_pipeline_path( + conf: 'Config', + gc: GirderClient, + pipeline: dict, + trained_pipeline_path: Path, +) -> Path: + if pipeline["type"] == constants.TrainedPipelineCategory: + gc.downloadFolderRecursive(pipeline["folderId"], str(trained_pipeline_path)) + return trained_pipeline_path / pipeline["pipe"] + return conf.get_extracted_pipeline_path() / pipeline["pipe"] + + +def _append_frame_range_video_settings( + command: List[str], + input_folder: GirderModel, + frame_range: Tuple[int, int], + pipeline_pipe: str, +) -> None: + command.append(f"-s downsampler:start_frame={shlex.quote(str(frame_range[0]))}") + command.append(f"-s downsampler:end_frame={shlex.quote(str(frame_range[1]))}") + input_fps = fromMeta(input_folder, constants.FPSMarker) + original_fps = fromMeta(input_folder, constants.OriginalFPSMarker, default=None) + is_native = original_fps is None or input_fps >= original_fps + command.append(f"-s downsampler:frame_range_is_native={str(is_native).lower()}") + renumber = pipeline_renumbers_frames(pipeline_pipe) + command.append(f"-s downsampler:renumber_frames={str(renumber).lower()}") + command.append(f"-s downsampler:adjust_timestamps={str(renumber).lower()}") + + +def _push_new_dataset_from_media( + gc: GirderClient, + manager: JobManager, + params: PipelineJob, + input_folder_id: str, + output_path: Path, + pipeline: dict, + *, + transcoded_video: Optional[str] = None, +) -> None: + """Create a sibling dataset from KWIVER media output (filter/transcode/disparity).""" + output_dataset_name = params.get('output_dataset_name') or ( + f"{pipeline.get('name', 'pipeline')}_output" + ) + output_parent_folder_id = params.get('output_parent_folder_id') + input_folder = gc.getFolder(input_folder_id) + source_fps = fromMeta(input_folder, constants.FPSMarker, default=-1) + + if is_transcode_pipeline(pipeline) and transcoded_video: + # Prefer uploading only the produced video via a dedicated staging dir + # when other files may be present under output_path. + staging = output_path / '_dataset_media' + utils.make_directory(staging) + video_path = Path(transcoded_video) + if not video_path.exists(): + # Fallback: first mp4 under output_path + videos = sorted(output_path.glob('*.mp4')) + if not videos: + raise Exception('Transcode pipeline produced no video file') + video_path = videos[0] + staged = staging / video_path.name + if video_path.resolve() != staged.resolve(): + shutil.copy2(video_path, staged) + utils.create_sibling_dataset_from_media( + gc, + manager, + input_folder_id, + staging, + output_dataset_name, + constants.VideoType, + source_fps, + parent_folder_id=output_parent_folder_id, + ) + return + + utils.create_sibling_dataset_from_media( + gc, + manager, + input_folder_id, + output_path, + output_dataset_name, + constants.ImageSequenceType, + source_fps, + parent_folder_id=output_parent_folder_id, + ) + + +def _inject_dataset_metadata_file(command, gc, working_dir: Path, params, manager) -> None: + """ + Download the dataset's optional metadata file (if the pipeline opted in) and + append its `-s =` override. Shared by the single and multicam + command-building branches. + """ + metadata_file_item_id = params.get('metadata_file_item_id') + metadata_file_key = params.get('metadata_file_key') + if not (metadata_file_item_id and metadata_file_key): + return + md_item = gc.getItem(metadata_file_item_id) + md_dir = utils.make_directory(working_dir / 'metadata_file') + gc.downloadItem(metadata_file_item_id, str(md_dir), name=md_item.get('name')) + # Locate what actually landed rather than reconstructing md_dir/: girder_client + # nests the download under a directory of that name when the item's file is named differently + # from the item (a sidecar renamed after upload), and it sanitizes the name with + # transformFilename first. Both make the reconstructed path wrong -- and in the nested case it + # is a directory, so an exists() check passes and binds a directory into the KWIVER setting. + # md_dir is created fresh for this item, so anything under it is its content. + downloaded = next((path for path in sorted(md_dir.rglob('*')) if path.is_file()), None) + if downloaded is not None: + append_metadata_file_kwiver_settings(command, downloaded, metadata_file_key) + else: + manager.write( + f'Warning: metadata item {metadata_file_item_id} ' + f'has no downloadable file under {md_dir}\n' + ) + + +def _append_input_list_kwiver_settings(command, pipeline, image_lists) -> None: + """ + Bind the run's per-camera input image lists to the KWIVER keys a pipe declares + via `# Image List Keys:`. image_lists is one single-file, line-separated list + per camera. A key template containing `{cam}` is expanded per camera (1-based) + — e.g. `stabilizer:image_list{cam}` -> image_list1, image_list2, ...; a key + without `{cam}` gets the first camera's list. Sea-lion registration needs the + list here in addition to the input reader's video_filename. + """ + if not image_lists: + return + for key in (pipeline.get('metadata') or {}).get('imageListKeys') or []: + if '{cam}' in key: + for idx, image_list in enumerate(image_lists, start=1): + expanded = key.replace('{cam}', str(idx)) + command.append(f'-s {shlex.quote(expanded)}={shlex.quote(image_list)}') + else: + command.append(f'-s {shlex.quote(key)}={shlex.quote(image_lists[0])}') + + +def _find_stereo_calibration_outputs(output_dir: Path) -> List[Path]: + """Return likely calibration outputs written by stereo calibration pipelines.""" + candidates: List[Path] = [] + for path in output_dir.iterdir(): + if not path.is_file(): + continue + lower_name = path.name.lower() + if 'calibration' not in lower_name: + continue + if not constants.stereoCalibrationRegex.search(path.name): + continue + candidates.append(path) + return sorted(candidates, key=lambda p: p.name.lower()) + + +@app.task(bind=True, acks_late=True, ignore_result=True) +def run_pipeline(self: Task, params: PipelineJob): + conf = Config() + conf.require_viame_install() + context: dict = {} + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + gc: GirderClient = self.girder_client + utils.authenticate_urllib(gc) + manager.updateStatus(JobStatus.FETCHING_INPUT) + + # Extract params + pipeline = params["pipeline"] + input_folder_id = str(params["input_folder"]) + input_type = params["input_type"] + output_folder_id = str(params["output_folder"]) + input_revision = params["input_revision"] + force_transcoded = params.get('force_transcoded', False) + runtime_params = params.get('runtime_params') or {} + frame_range = runtime_params.get('frameRange') + multicam_params: MulticamPipelineJob = params + multicam_cameras: List[MulticamCameraJob] = multicam_params.get('multicam_cameras') or [] + camera_name = params.get('camera_name') + if camera_name: + # Log non-default camera targets so job history shows which view ran. + default_display = multicam_params.get('multicam_default_display') + if not default_display or camera_name != default_display: + print(f'Running pipeline on camera: {camera_name}') + with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): + _working_directory_path = Path(_working_directory) + input_path = utils.make_directory(_working_directory_path / 'input') + trained_pipeline_path = utils.make_directory(_working_directory_path / 'trained_pipeline') + output_path = utils.make_directory(_working_directory_path / 'output') + + detector_output_file = str(output_path / 'detector_output.csv') + track_output_file = str(output_path / 'track_output.csv') + img_list_path = input_path / 'img_list_file.txt' + + pipeline_path = _resolve_pipeline_path(conf, gc, pipeline, trained_pipeline_path) + + assert pipeline_path.exists(), ( + "Requested pipeline could not be found." + " Make sure that VIAME is installed correctly and all addons have loaded." + f" Job asked for {pipeline_path} but it does not exist" + ) + + if multicam_cameras: + input_folder = gc.getFolder(input_folder_id) + input_fps = fromMeta(input_folder, constants.FPSMarker) + requires_input = multicam_params.get('multicam_requires_input', False) + creates_new_dataset = pipeline_creates_new_dataset(pipeline) + camera_media: Dict[str, Tuple[List[str], str]] = {} + + for cam_index, camera in enumerate(multicam_cameras, start=1): + cam_input_path = utils.make_directory(input_path / camera['name']) + media_list, media_type = utils.download_source_media( + gc, camera['folder_id'], cam_input_path, force_transcoded + ) + if frame_range is not None and media_type == constants.ImageSequenceType: + media_list = filter_image_list_by_frame_range(media_list, frame_range) + camera_media[camera['name']] = (media_list, media_type) + if requires_input and camera.get('input_revision') is not None: + gt_path = _working_directory_path / f'detections{cam_index}.csv' + utils.download_revision_csv( + gc, camera['folder_id'], camera['input_revision'], gt_path + ) + + arg_file_pair, out_files = build_multicam_kwiver_settings( + _working_directory_path, + multicam_cameras, + camera_media, + requires_input=requires_input, + ) + + command = [ + f". {shlex.quote(str(conf.viame_setup_script))} &&", + f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", + "viame runner", + f"-p {shlex.quote(str(pipeline_path))}", + ] + if input_type == constants.VideoType: + command.extend( + [ + '-s input:video_reader:type=vidl_ffmpeg', + f"-s downsampler:target_frame_rate={shlex.quote(str(input_fps))}", + ] + ) + if frame_range is not None: + _append_frame_range_video_settings( + command, input_folder, frame_range, pipeline['pipe'] + ) + for arg, file_name in arg_file_pair.items(): + command.append(f"-s {shlex.quote(arg)}={shlex.quote(file_name)}") + + transcoded_video: Optional[str] = None + if creates_new_dataset: + video_name = None + if is_transcode_pipeline(pipeline): + video_name = str( + output_path / f"{pipeline.get('name', 'transcode')}_{input_folder_id}.mp4" + ) + transcoded_video = append_new_dataset_media_writers( + command, pipeline, output_path, video_filename=video_name + ) + + calibration_item_id = multicam_params.get('calibration_item_id') + if calibration_item_id and is_stereo_measurement_pipeline(pipeline): + cal_item = gc.getItem(calibration_item_id) + cal_dir = utils.make_directory(_working_directory_path / 'calibration') + gc.downloadItem( + calibration_item_id, + str(cal_dir), + name=cal_item.get('name'), + ) + cal_path = find_downloaded_calibration_file(cal_dir) + if cal_path is not None: + append_stereo_calibration_kwiver_settings(command, cal_path, pipeline) + else: + manager.write( + f'Warning: calibration item {calibration_item_id} ' + f'has no recognized calibration file under {cal_dir}\n' + ) + + # One image list per camera (each a single line-separated file). + input_manifests = [ + arg_file_pair[f'input{i + 1}:video_filename'] + for i in range(len(multicam_cameras)) + if f'input{i + 1}:video_filename' in arg_file_pair + ] + _append_input_list_kwiver_settings(command, pipeline, input_manifests) + + _inject_dataset_metadata_file(command, gc, _working_directory_path, params, manager) + + kwiver_params = params.get('kwiver_params') + if kwiver_params: + for key, value in kwiver_params.items(): + command.append(f'-s {shlex.quote(key)}={shlex.quote(str(value))}') + + manager.updateStatus(JobStatus.RUNNING) + popen_kwargs = { + 'args': " ".join(command), + 'shell': True, + 'executable': '/bin/bash', + 'cwd': output_path, + 'env': conf.gpu_process_env, + } + utils.stream_subprocess(self, context, manager, popen_kwargs) + + if ( + is_stereo_measurement_pipeline(pipeline) + and 'calibrate_cameras' in str(pipeline.get('pipe', '')).lower() + ): + calibration_outputs = _find_stereo_calibration_outputs(output_path) + if calibration_outputs: + calibration_output = calibration_outputs[0] + try: + uploaded_calibration = gc.uploadFileToFolder( + input_folder_id, + str(calibration_output), + ) + uploaded_calibration_file_id = uploaded_calibration.get('_id') + if uploaded_calibration_file_id is not None: + uploaded_calibration_file_id_str = str(uploaded_calibration_file_id) + cal_url = ( + f'/dive_dataset/{input_folder_id}/calibration' + f'?fileId={uploaded_calibration_file_id_str}' + ) + gc.sendRestRequest('POST', cal_url) + manager.write( + 'Assigned calibration output to dataset: ' + f'{calibration_output.name}\n' + ) + else: + manager.write( + 'Warning: uploaded calibration output ' + f'{calibration_output.name} has no file id\n' + ) + except Exception as exc: + manager.write( + 'Warning: failed to assign calibration output ' + f'{calibration_output.name}: {exc}\n' + ) + else: + manager.write( + 'Warning: stereo calibration pipeline produced no ' + 'recognized calibration output file\n' + ) + + if creates_new_dataset: + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + _push_new_dataset_from_media( + gc, + manager, + params, + input_folder_id, + output_path, + pipeline, + transcoded_video=transcoded_video, + ) + return + + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + for camera in multicam_cameras: + cam_name = camera['name'] + output_name = out_files[cam_name] + # Multicam KWIVER args use basename-only writers; viame cwd is output_path, + # so CSVs are created under output/, not the temp directory root. + output_file = output_path / output_name + if not output_file.exists() or not output_file.stat().st_size: + detector_name = output_name.replace('computed_tracks', 'computed_detections') + detector_path = output_path / detector_name + if detector_path.exists() and detector_path.stat().st_size: + output_file = detector_path + if frame_range is not None and camera_media[cam_name][1] == constants.VideoType: + filtered_path = filter_csv_by_frame_range(str(output_file), frame_range) + output_file = Path(filtered_path) + newfile = gc.uploadFileToFolder(camera['folder_id'], str(output_file)) + gc.addMetadataToItem(str(newfile["itemId"]), {"pipeline": pipeline}) + gc.post( + f'dive_rpc/postprocess/{camera["folder_id"]}', + data={"skipJobs": True}, + ) + return + + # Download source media + input_folder: GirderModel = gc.getFolder(input_folder_id) + creates_new_dataset = pipeline_creates_new_dataset(pipeline) + input_media_list, _ = utils.download_source_media( + gc, input_folder_id, input_path, force_transcoded + ) + + if input_type == constants.VideoType: + input_fps = fromMeta(input_folder, constants.FPSMarker) + assert len(input_media_list) == 1, "Expected exactly 1 video" + command = [ + f". {shlex.quote(str(conf.viame_setup_script))} &&", + f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", + "viame runner", + "-s input:video_reader:type=vidl_ffmpeg", + f"-p {shlex.quote(str(pipeline_path))}", + f"-s input:video_filename={shlex.quote(input_media_list[0])}", + f"-s downsampler:target_frame_rate={shlex.quote(str(input_fps))}", + f"-s detector_writer:file_name={shlex.quote(detector_output_file)}", + f"-s track_writer:file_name={shlex.quote(track_output_file)}", + ] + if frame_range is not None: + _append_frame_range_video_settings( + command, input_folder, frame_range, pipeline['pipe'] + ) + elif input_type == constants.ImageSequenceType: + # Filter image list by frame range if specified + filtered_media_list = input_media_list + if frame_range is not None: + filtered_media_list = filter_image_list_by_frame_range( + input_media_list, frame_range + ) + with open(img_list_path, "w+") as img_list_file: + img_list_file.write('\n'.join(filtered_media_list)) + command = [ + f". {shlex.quote(str(conf.viame_setup_script))} &&", + f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", + "viame runner", + f"-p {shlex.quote(str(pipeline_path))}", + f"-s input:video_filename={shlex.quote(str(img_list_path))}", + f"-s detector_writer:file_name={shlex.quote(detector_output_file)}", + f"-s track_writer:file_name={shlex.quote(track_output_file)}", + ] + else: + raise ValueError('Unknown input type: {}'.format(input_type)) + + # Include input detections + if input_revision is not None: + pipeline_input_file = input_path / 'groundtruth.csv' + utils.download_revision_csv(gc, input_folder_id, input_revision, pipeline_input_file) + quoted_input_file = shlex.quote(str(pipeline_input_file)) + command.append(f'-s detection_reader:file_name={quoted_input_file}') + command.append(f'-s track_reader:file_name={quoted_input_file}') + + transcoded_video = None + if creates_new_dataset: + video_name = None + if is_transcode_pipeline(pipeline): + video_name = str( + output_path / f"{pipeline.get('name', 'transcode')}_{input_folder_id}.mp4" + ) + transcoded_video = append_new_dataset_media_writers( + command, pipeline, output_path, video_filename=video_name + ) + + single_input_manifest = ( + str(img_list_path) if input_type == constants.ImageSequenceType else input_media_list[0] + ) + _append_input_list_kwiver_settings(command, pipeline, [single_input_manifest]) + + _inject_dataset_metadata_file(command, gc, _working_directory_path, params, manager) + + # Apply user-provided KWIVER parameter overrides. + kwiver_params = params.get('kwiver_params') + if kwiver_params: + for key, value in kwiver_params.items(): + command.append(f'-s {shlex.quote(key)}={shlex.quote(str(value))}') + + manager.updateStatus(JobStatus.RUNNING) + popen_kwargs = { + 'args': " ".join(command), + 'shell': True, + 'executable': '/bin/bash', + 'cwd': output_path, + 'env': conf.gpu_process_env, + } + utils.stream_subprocess(self, context, manager, popen_kwargs) + + if creates_new_dataset: + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + _push_new_dataset_from_media( + gc, + manager, + params, + input_folder_id, + output_path, + pipeline, + transcoded_video=transcoded_video, + ) + return + + if Path(track_output_file).exists() and os.path.getsize(track_output_file): + output_file = track_output_file + else: + output_file = detector_output_file + + # Filter output CSV by frame range for videos + if frame_range is not None and input_type == constants.VideoType: + output_file = filter_csv_by_frame_range(output_file, frame_range) + + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + newfile = gc.uploadFileToFolder(output_folder_id, output_file) + + gc.addMetadataToItem(str(newfile["itemId"]), {"pipeline": pipeline}) + gc.post(f'dive_rpc/postprocess/{output_folder_id}', data={"skipJobs": True}) diff --git a/server/dive_tasks/run_training.py b/server/dive_tasks/run_training.py new file mode 100644 index 000000000..013a3b47f --- /dev/null +++ b/server/dive_tasks/run_training.py @@ -0,0 +1,207 @@ +from contextlib import suppress +from pathlib import Path +import shlex +import tempfile +from typing import List, Tuple + +from girder_client import GirderClient +from girder_worker.app import app +from girder_worker.task import Task +from girder_worker.utils import JobManager, JobStatus + +from dive_tasks import utils +from dive_tasks.manager import patch_manager +from dive_tasks.viame_config import Config +from dive_utils import constants +from dive_utils.types import ExportTrainedPipelineJob, TrainingJob + + +@app.task(bind=True, acks_late=True, ignore_results=True) +def export_trained_pipeline(self: Task, params: ExportTrainedPipelineJob): + conf = Config() + conf.require_viame_install() + context: dict = {} + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + gc: GirderClient = self.girder_client + utils.authenticate_urllib(gc) + manager.updateStatus(JobStatus.FETCHING_INPUT) + + # Extract params + input_folder_id = params["input_folder"] + output_folder_id = params["output_folder"] + output_name = params["output_name"] + + with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): + _working_directory_path = Path(_working_directory) + trained_pipeline_path = utils.make_directory(_working_directory_path / 'trained_pipeline') + output_path = utils.make_directory(_working_directory_path / 'output') + onnx_path = output_path / output_name + convert_to_onnx_pipeline_path = conf.viame_pipeline_path / "convert_model_to_onnx.pipe" + + gc.downloadFolderRecursive(input_folder_id, str(trained_pipeline_path)) + extensions = ['*.weights', '*.ckpt', '*.pth'] + model_file = None + + for ext in extensions: + found_files = list(trained_pipeline_path.glob(ext)) + if found_files: + model_file = found_files[0] + break + + if not model_file: + raise FileNotFoundError(f"No weights path ({extensions}) found.") + + # Convert pipeline to ONNX + command = [ + f". {shlex.quote(str(conf.viame_setup_script))} &&", + f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", + "viame runner", + f"-p {shlex.quote(str(convert_to_onnx_pipeline_path))}", + f"-s onnx_convert:model_path={shlex.quote(str(model_file))}", + f"-s onnx_convert:onnx_model_prefix={shlex.quote(str(onnx_path))}", + ] + + manager.updateStatus(JobStatus.RUNNING) + popen_kwargs = { + 'args': " ".join(command), + 'shell': True, + 'executable': '/bin/bash', + 'cwd': output_path, + 'env': conf.gpu_process_env, + } + utils.stream_subprocess(self, context, manager, popen_kwargs) + + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + gc.uploadFileToFolder(output_folder_id, onnx_path) + + +@app.task(bind=True, acks_late=True, ignore_result=True) +def train_pipeline(self: Task, params: TrainingJob): + """Train a pipeline by making a call to viame train""" + conf = Config() + conf.require_viame_install() + context: dict = {} + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + gc: GirderClient = self.girder_client + utils.authenticate_urllib(gc) + manager.updateStatus(JobStatus.FETCHING_INPUT) + + # Extract params + results_folder_id = params['results_folder_id'] + dataset_input_list = params['dataset_input_list'] + pipeline_name = params['pipeline_name'] + config = params['config'] + annotated_frames_only = params['annotated_frames_only'] + label_text = params['label_txt'] + model = params.get('model', None) + # Normalize: model can arrive as a list of [key, value] pairs from some serialization paths + if model is not None and isinstance(model, list): + model = dict(model) + force_transcoded = params.get('force_transcoded', False) + + pipeline_base_path = Path(conf.get_extracted_pipeline_path()) + config_file = pipeline_base_path / config + # List of (input folder, ground truth file) pairs for creating input lists + input_groundtruth_list: List[Tuple[Path, Path]] = [] + # root_data_dir is the directory passed to `viame train` + with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): + _working_directory_path = Path(_working_directory) + input_path = utils.make_directory(_working_directory_path / 'input') + output_path = utils.make_directory(_working_directory_path / 'output') + + for source_folder_id, revision in dataset_input_list: + download_path = utils.make_directory(input_path / source_folder_id) + groundtruth_path = download_path / 'groundtruth.csv' + # Download groundtruth item + utils.download_revision_csv(gc, source_folder_id, revision, groundtruth_path) + # Download input media + input_media_list, input_type = utils.download_source_media( + gc, source_folder_id, download_path, force_transcoded + ) + if input_type == constants.VideoType: + download_path = Path(input_media_list[0]) + # Set media source location + input_groundtruth_list.append((download_path, groundtruth_path)) + + input_folder_file_list = input_path / "input_folder_list.txt" + ground_truth_file_list = input_path / "input_truth_list.txt" + with open(input_folder_file_list, "w+") as data_list: + with open(ground_truth_file_list, "w+") as truth_list: + for folder_path, groundtruth_path in input_groundtruth_list: + data_list.write(f"{folder_path}\n") + truth_list.write(f"{groundtruth_path}\n") + + training_results_path = utils.make_directory(output_path / "category_models") + + command = [ + f". {shlex.quote(str(conf.viame_setup_script))} &&", + f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", + f"{shlex.quote(str(conf.viame_executable))} train", + "--input-list", + shlex.quote(str(input_folder_file_list)), + "--input-truth", + shlex.quote(str(ground_truth_file_list)), + "--config", + shlex.quote(str(config_file)), + "--no-query", + "--no-embedded-pipe", + ] + + if annotated_frames_only: + command.append("--gt-frames-only") + + if label_text: + labels_path = input_path / "labels.txt" + with open(labels_path, "w+") as labels_file: + labels_file.write(label_text) + command.append("--labels") + command.append(shlex.quote(str(labels_path))) + + if model: + model_path = None + if model.get('folderId', False): + trained_pipeline_path = utils.make_directory( + _working_directory_path / 'trained_pipeline' + ) + gc.downloadFolderRecursive(model["folderId"], str(trained_pipeline_path)) + model_path = trained_pipeline_path / model["name"] + elif model.get('path', False): + model_path = model['path'] + if model_path: + command.append("--init-weights") + command.append(shlex.quote(str(model_path))) + + manager.updateStatus(JobStatus.RUNNING) + popen_kwargs = { + 'args': " ".join(command), + 'shell': True, + 'executable': '/bin/bash', + 'cwd': output_path, + 'env': conf.gpu_process_env, + } + utils.stream_subprocess(self, context, manager, popen_kwargs) + + # Check that there are results in the output path + if len(list(training_results_path.glob("*"))) == 0: + raise RuntimeError("Training output didn't produce results, discarding...") + + manager.updateStatus(JobStatus.PUSHING_OUTPUT) + # This is the name of the folder that is uploaded to the + # "Training Results" girder folder + girder_output_folder = gc.createFolder( + results_folder_id, + pipeline_name, + metadata={ + constants.TrainedPipelineMarker: True, + "trained_on": dataset_input_list, + }, + ) + gc.upload(f"{training_results_path}/*", girder_output_folder["_id"]) diff --git a/server/dive_tasks/tasks.py b/server/dive_tasks/tasks.py index c5f7b587c..5e8762943 100644 --- a/server/dive_tasks/tasks.py +++ b/server/dive_tasks/tasks.py @@ -1,1612 +1,52 @@ -from contextlib import suppress -import logging -import os -from pathlib import Path -import shlex -import shutil -import tempfile -from typing import Dict, List, Optional, Tuple -from urllib import request -from urllib.parse import urlparse -import zipfile - -from GPUtil import getGPUs -import gdown -from gdown.parse_url import is_google_drive_url, parse_url -from girder_client import GirderClient, HttpError -from girder_worker.app import app -from girder_worker.task import Task -from girder_worker.utils import JobManager, JobStatus - -from dive_tasks import utils -from dive_tasks.frame_alignment import check_and_fix_frame_alignment, is_frame_misaligned -from dive_tasks.manager import patch_manager -from dive_tasks.multicam_pipeline import ( - append_metadata_file_kwiver_settings, - append_stereo_calibration_kwiver_settings, - build_multicam_kwiver_settings, - find_downloaded_calibration_file, - is_stereo_measurement_pipeline, +"""Compatibility barrel re-exporting Celery tasks and helpers. + +Prefer importing from the focused modules directly in new code: +``convert_video``, ``convert_images``, ``run_pipeline``, ``run_training``, +``upgrade_pipelines``, and ``viame_config``. +""" + +from dive_tasks.convert_images import ( + convert_calibration, + convert_images, + convert_large_images, + extract_zip, ) -from dive_tasks.pipeline_creates_dataset import ( - append_new_dataset_media_writers, - is_transcode_pipeline, - pipeline_creates_new_dataset, - pipeline_renumbers_frames, +from dive_tasks.convert_video import convert_video, resolve_annotation_fps +from dive_tasks.run_pipeline import ( + _inject_dataset_metadata_file, + filter_csv_by_frame_range, + filter_image_list_by_frame_range, + run_pipeline, ) -from dive_tasks.pipeline_discovery import discover_configs -from dive_utils import asbool, calibration_format, constants, fromMeta -from dive_utils.types import ( - AvailableJobSchema, - ExportTrainedPipelineJob, - GirderModel, - MulticamCameraJob, - MulticamPipelineJob, - PipelineJob, - TrainingJob, +from dive_tasks.run_training import export_trained_pipeline, train_pipeline +from dive_tasks.upgrade_pipelines import ( + UPGRADE_JOB_DEFAULT_URLS, + _addon_zip_path_for_url, + download_google_drive_zip, + is_google_drive_addon_url, + upgrade_pipelines, ) - -logger = logging.getLogger(__name__) - - -def filter_csv_by_frame_range(csv_path: str, frame_range: Tuple[int, int]) -> str: - """Filter VIAME CSV to only include detections within frame range. - - Args: - csv_path: Path to the input CSV file - frame_range: Tuple of (start_frame, end_frame) inclusive - - Returns: - Path to the filtered CSV file - """ - start_frame, end_frame = frame_range - filtered_path = csv_path.replace('.csv', '_filtered.csv') - - with open(csv_path, 'r') as infile, open(filtered_path, 'w') as outfile: - for line in infile: - if line.startswith('#'): - outfile.write(line) - continue - parts = line.split(',') - if len(parts) >= 3: - try: - frame = int(parts[2]) # Frame number is column 3 (0-indexed as column 2) - if start_frame <= frame <= end_frame: - outfile.write(line) - except ValueError: - # If frame number can't be parsed, include the line - outfile.write(line) - return filtered_path - - -def filter_image_list_by_frame_range( - image_list: List[str], frame_range: Tuple[int, int] -) -> List[str]: - """Filter an image list to only include images within frame range. - - Args: - image_list: List of image file paths - frame_range: Tuple of (start_frame, end_frame) inclusive (0-indexed) - - Returns: - Filtered list of image file paths - """ - start_frame, end_frame = frame_range - # Ensure we don't go out of bounds - start_frame = max(0, start_frame) - end_frame = min(end_frame, len(image_list) - 1) - return image_list[start_frame : end_frame + 1] - - -EMPTY_JOB_SCHEMA: AvailableJobSchema = { - 'pipelines': {}, - 'training': { - 'configs': [], - 'default': None, - }, - 'models': {}, -} - -# https://github.com/VIAME/VIAME/blob/master/cmake/download_viame_addons.csv -UPGRADE_JOB_DEFAULT_URLS: List[str] = [ - 'https://viame.kitware.com/api/v1/item/627b145487bad2e19a4c4697/download', # HabCam - 'https://viame.kitware.com/api/v1/item/627b32b1994809b024f207a7/download', # SEFSC - 'https://viame.kitware.com/api/v1/item/627b3289ea630db5587b577d/download', # SWFSC-PengHead - 'https://viame.kitware.com/api/v1/item/627b326fea630db5587b577b/download', # Motion - 'https://viame.kitware.com/api/v1/item/627b326cc4da86e2cd3abb5b/download', # EM Tuna - 'https://viame.kitware.com/api/v1/item/627b3282c4da86e2cd3abb5d/download', # MOUSS - 'https://viame.kitware.com/api/v1/item/615bc7aa7e5c13a5bb9af7a7/download', # Aerial Penguin - 'https://viame.kitware.com/api/v1/item/629807c192adc2f0ecfa5b54/download', # Sea Lion +from dive_tasks.viame_config import EMPTY_JOB_SCHEMA, Config, get_gpu_environment + +__all__ = [ + 'Config', + 'EMPTY_JOB_SCHEMA', + 'UPGRADE_JOB_DEFAULT_URLS', + '_addon_zip_path_for_url', + '_inject_dataset_metadata_file', + 'convert_calibration', + 'convert_images', + 'convert_large_images', + 'convert_video', + 'download_google_drive_zip', + 'export_trained_pipeline', + 'extract_zip', + 'filter_csv_by_frame_range', + 'filter_image_list_by_frame_range', + 'get_gpu_environment', + 'is_google_drive_addon_url', + 'resolve_annotation_fps', + 'run_pipeline', + 'train_pipeline', + 'upgrade_pipelines', ] - - -def get_gpu_environment() -> Dict[str, str]: - """Get environment variables for using CUDA enabled GPUs.""" - env = os.environ.copy() - - gpu_uuid = env.get("WORKER_GPU_UUID") - gpus = [gpu.id for gpu in getGPUs() if gpu.uuid == gpu_uuid] - - # Only set this env var if WORKER_GPU_UUID was supplied, - # and it matches an installed GPU - if gpus: - env["CUDA_VISIBLE_DEVICES"] = str(gpus[0]) - # Support for NOAA python3.10 means removing the local venv from the path - env["PATH"] = env.get("PATH").replace("/opt/dive/local/venv/bin", "") - return env - - -_VIAME_WORKER_QUEUES = frozenset({'pipelines', 'training'}) - - -def _worker_requires_viame_install() -> bool: - """ - Only pipeline/training workers need a local VIAME install. - - Default (``celery``), ``local``, and dev ``localworker`` processes do not; - they never call :class:`Config` today, but this keeps :meth:`Config.__init__` - safe if a task is misrouted. - """ - queues = os.environ.get('WORKER_WATCHING_QUEUES', '') - watched = {q.strip() for q in queues.split(',') if q.strip()} - return bool(watched & _VIAME_WORKER_QUEUES) - - -class Config: - def __init__(self): - self.gpu_process_env = get_gpu_environment() - self.viame_install_directory = os.environ.get( - 'VIAME_INSTALL_PATH', - '/opt/noaa/viame', - ) - self.addon_root_directory = os.environ.get( - 'ADDON_ROOT_DIR', - '/tmp/addons', - ) - self.kwiver_log_level = os.environ.get( - 'KWIVER_DEFAULT_LOG_LEVEL', - 'warn', - ) - - self.pipeline_subdir = 'configs/pipelines' - self.viame_install_path = Path(self.viame_install_directory) - self.viame_setup_script = self.viame_install_path / "setup_viame.sh" - self.viame_executable = self.viame_install_path / "bin" / "viame" - self.viame_pipeline_path = self.viame_install_path / self.pipeline_subdir - - if _worker_requires_viame_install(): - self.require_viame_install() - - self.addon_root_path = Path(self.addon_root_directory) - self.addon_zip_path = utils.make_directory(self.addon_root_path / 'zips') - self.addon_extracted_path = utils.make_directory(self.addon_root_path / 'extracted') - - # Set include directory to include pipelines from this path - # https://github.com/VIAME/VIAME/issues/131 - self.gpu_process_env['SPROKIT_PIPE_INCLUDE_PATH'] = str( - self.addon_extracted_path / self.pipeline_subdir - ) - - def require_viame_install(self) -> None: - assert self.viame_install_path.exists(), "VIAME Base install directory missing." - assert self.viame_setup_script.is_file(), "VIAME Setup Script missing" - assert self.viame_executable.is_file(), "VIAME Executable missing" - assert self.viame_pipeline_path.exists(), "VIAME common pipe directory missing." - - def get_extracted_pipeline_path(self, missing_ok=False) -> Path: - """ - Includes subdirectory for pipelines - """ - pipeline_path = self.addon_extracted_path / self.pipeline_subdir - if not missing_ok: - assert pipeline_path.exists(), f"Missing path {pipeline_path}" - return pipeline_path - - -def _normalize_google_drive_url(url: str) -> str: - """Strip a leading www. so gdown recognizes common pasted Drive links.""" - parsed = urlparse(url) - host = parsed.netloc.lower() - if host.startswith('www.'): - return parsed._replace(netloc=host[4:]).geturl() - return url - - -def is_google_drive_addon_url(url: str) -> bool: - """Return True if url is a Google Drive link (after normalizing www.).""" - return is_google_drive_url(_normalize_google_drive_url(url)) - - -def download_google_drive_zip(url: str, dest: Path) -> None: - """Download a publicly shared Google Drive zip to dest via gdown.""" - gdown.download(url=_normalize_google_drive_url(url), output=str(dest), quiet=True) - - -def _addon_zip_path_for_url(addon_url: str, addon_zip_dir: Path) -> Path: - normalized = _normalize_google_drive_url(addon_url) - if is_google_drive_url(normalized): - file_id, _ = parse_url(normalized) - if file_id: - return addon_zip_dir / f'gdrive_{file_id}.zip' - download_name = urlparse(addon_url).path.replace(os.path.sep, '_') - return addon_zip_dir / f'{download_name}.zip' - - -@app.task(bind=True, acks_late=True, ignore_result=True) -def upgrade_pipelines( - self: Task, - urls: List[str] = UPGRADE_JOB_DEFAULT_URLS, - force: bool = False, -): - """Install addons from zip files over HTTP (including Google Drive share links)""" - conf = Config() - context: dict = {} - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - gc: GirderClient = self.girder_client - # zipfiles to extract after download is complete - addons_to_update_update: List[Path] = [] - - for addon in urls: - zipfile_path = _addon_zip_path_for_url(addon, conf.addon_zip_path) - had_existing_zip = zipfile_path.exists() - try: - if not had_existing_zip or force: - manager.write(f'Downloading {addon} to {zipfile_path}\n') - if is_google_drive_addon_url(addon): - download_google_drive_zip(addon, zipfile_path) - else: - request.urlretrieve(addon, filename=zipfile_path) - else: - manager.write(f'Skipping download of {zipfile_path}\n') - addons_to_update_update.append(zipfile_path) - except Exception as exc: - logger.exception('Failed to download addon %s', addon) - manager.write(f'Failed to download {addon}: {exc}\nSkipping.\n') - if zipfile_path.exists() and not had_existing_zip: - zipfile_path.unlink(missing_ok=True) - if utils.check_canceled(self, context, force=False): - manager.updateStatus(JobStatus.CANCELED) - return - - # remove and recreate the existing addon pipeline directory - shutil.rmtree(conf.addon_extracted_path) - # Seed base pipelines from the VIAME image when available (GPU workers only). - if conf.viame_pipeline_path.exists(): - shutil.copytree(conf.viame_pipeline_path, conf.get_extracted_pipeline_path(missing_ok=True)) - # Extract zipfiles over newly copied files. Right now the zip archives - # MUST contain the pipeline subdir (e.g. configs/pipelines) in their - # internal structure. - for zipfile_path in addons_to_update_update: - manager.write(f'Extracting {zipfile_path} to {str(conf.addon_extracted_path)}\n') - z = zipfile.ZipFile(zipfile_path) - z.extractall(conf.addon_extracted_path) - - if utils.check_canceled(self, context): - # Remove everything - shutil.rmtree(conf.addon_extracted_path) - manager.updateStatus(JobStatus.CANCELED) - gc.put('dive_configuration/static_pipeline_configs', json=EMPTY_JOB_SCHEMA) - return - - # finally, crawl the new files and report results - summary = discover_configs(conf.get_extracted_pipeline_path()) - manager.write(str(summary)) - gc.put('dive_configuration/static_pipeline_configs', json=summary) - # get a list of files in the zip directory for the installed configuration listing - downloaded = [] - - # Iterate directory - for path in os.listdir(conf.addon_zip_path): - # check if current path is a file - if os.path.isfile(os.path.join(conf.addon_zip_path, path)): - downloaded.append(path) - print('Downloaded Files') - print(downloaded) - gc.put('dive_configuration/installed_addons', json={'downloaded': downloaded}) - - -def _resolve_pipeline_path( - conf: 'Config', - gc: GirderClient, - pipeline: dict, - trained_pipeline_path: Path, -) -> Path: - if pipeline["type"] == constants.TrainedPipelineCategory: - gc.downloadFolderRecursive(pipeline["folderId"], str(trained_pipeline_path)) - return trained_pipeline_path / pipeline["pipe"] - return conf.get_extracted_pipeline_path() / pipeline["pipe"] - - -def _append_frame_range_video_settings( - command: List[str], - input_folder: GirderModel, - frame_range: Tuple[int, int], - pipeline_pipe: str, -) -> None: - command.append(f"-s downsampler:start_frame={shlex.quote(str(frame_range[0]))}") - command.append(f"-s downsampler:end_frame={shlex.quote(str(frame_range[1]))}") - input_fps = fromMeta(input_folder, constants.FPSMarker) - original_fps = fromMeta(input_folder, constants.OriginalFPSMarker, default=None) - is_native = original_fps is None or input_fps >= original_fps - command.append(f"-s downsampler:frame_range_is_native={str(is_native).lower()}") - renumber = pipeline_renumbers_frames(pipeline_pipe) - command.append(f"-s downsampler:renumber_frames={str(renumber).lower()}") - command.append(f"-s downsampler:adjust_timestamps={str(renumber).lower()}") - - -def _push_new_dataset_from_media( - gc: GirderClient, - manager: JobManager, - params: PipelineJob, - input_folder_id: str, - output_path: Path, - pipeline: dict, - *, - transcoded_video: Optional[str] = None, -) -> None: - """Create a sibling dataset from KWIVER media output (filter/transcode/disparity).""" - output_dataset_name = params.get('output_dataset_name') or ( - f"{pipeline.get('name', 'pipeline')}_output" - ) - output_parent_folder_id = params.get('output_parent_folder_id') - input_folder = gc.getFolder(input_folder_id) - source_fps = fromMeta(input_folder, constants.FPSMarker, default=-1) - - if is_transcode_pipeline(pipeline) and transcoded_video: - # Prefer uploading only the produced video via a dedicated staging dir - # when other files may be present under output_path. - staging = output_path / '_dataset_media' - utils.make_directory(staging) - video_path = Path(transcoded_video) - if not video_path.exists(): - # Fallback: first mp4 under output_path - videos = sorted(output_path.glob('*.mp4')) - if not videos: - raise Exception('Transcode pipeline produced no video file') - video_path = videos[0] - staged = staging / video_path.name - if video_path.resolve() != staged.resolve(): - shutil.copy2(video_path, staged) - utils.create_sibling_dataset_from_media( - gc, - manager, - input_folder_id, - staging, - output_dataset_name, - constants.VideoType, - source_fps, - parent_folder_id=output_parent_folder_id, - ) - return - - utils.create_sibling_dataset_from_media( - gc, - manager, - input_folder_id, - output_path, - output_dataset_name, - constants.ImageSequenceType, - source_fps, - parent_folder_id=output_parent_folder_id, - ) - - -def _inject_dataset_metadata_file(command, gc, working_dir: Path, params, manager) -> None: - """ - Download the dataset's optional metadata file (if the pipeline opted in) and - append its `-s =` override. Shared by the single and multicam - command-building branches. - """ - metadata_file_item_id = params.get('metadata_file_item_id') - metadata_file_key = params.get('metadata_file_key') - if not (metadata_file_item_id and metadata_file_key): - return - md_item = gc.getItem(metadata_file_item_id) - md_dir = utils.make_directory(working_dir / 'metadata_file') - gc.downloadItem(metadata_file_item_id, str(md_dir), name=md_item.get('name')) - # Locate what actually landed rather than reconstructing md_dir/: girder_client - # nests the download under a directory of that name when the item's file is named differently - # from the item (a sidecar renamed after upload), and it sanitizes the name with - # transformFilename first. Both make the reconstructed path wrong -- and in the nested case it - # is a directory, so an exists() check passes and binds a directory into the KWIVER setting. - # md_dir is created fresh for this item, so anything under it is its content. - downloaded = next((path for path in sorted(md_dir.rglob('*')) if path.is_file()), None) - if downloaded is not None: - append_metadata_file_kwiver_settings(command, downloaded, metadata_file_key) - else: - manager.write( - f'Warning: metadata item {metadata_file_item_id} ' - f'has no downloadable file under {md_dir}\n' - ) - - -def _append_input_list_kwiver_settings(command, pipeline, image_lists) -> None: - """ - Bind the run's per-camera input image lists to the KWIVER keys a pipe declares - via `# Image List Keys:`. image_lists is one single-file, line-separated list - per camera. A key template containing `{cam}` is expanded per camera (1-based) - — e.g. `stabilizer:image_list{cam}` -> image_list1, image_list2, ...; a key - without `{cam}` gets the first camera's list. Sea-lion registration needs the - list here in addition to the input reader's video_filename. - """ - if not image_lists: - return - for key in (pipeline.get('metadata') or {}).get('imageListKeys') or []: - if '{cam}' in key: - for idx, image_list in enumerate(image_lists, start=1): - expanded = key.replace('{cam}', str(idx)) - command.append(f'-s {shlex.quote(expanded)}={shlex.quote(image_list)}') - else: - command.append(f'-s {shlex.quote(key)}={shlex.quote(image_lists[0])}') - - -def _find_stereo_calibration_outputs(output_dir: Path) -> List[Path]: - """Return likely calibration outputs written by stereo calibration pipelines.""" - candidates: List[Path] = [] - for path in output_dir.iterdir(): - if not path.is_file(): - continue - lower_name = path.name.lower() - if 'calibration' not in lower_name: - continue - if not constants.stereoCalibrationRegex.search(path.name): - continue - candidates.append(path) - return sorted(candidates, key=lambda p: p.name.lower()) - - -@app.task(bind=True, acks_late=True, ignore_result=True) -def run_pipeline(self: Task, params: PipelineJob): - conf = Config() - conf.require_viame_install() - context: dict = {} - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - gc: GirderClient = self.girder_client - utils.authenticate_urllib(gc) - manager.updateStatus(JobStatus.FETCHING_INPUT) - - # Extract params - pipeline = params["pipeline"] - input_folder_id = str(params["input_folder"]) - input_type = params["input_type"] - output_folder_id = str(params["output_folder"]) - input_revision = params["input_revision"] - force_transcoded = params.get('force_transcoded', False) - runtime_params = params.get('runtime_params') or {} - frame_range = runtime_params.get('frameRange') - multicam_params: MulticamPipelineJob = params - multicam_cameras: List[MulticamCameraJob] = multicam_params.get('multicam_cameras') or [] - camera_name = params.get('camera_name') - if camera_name: - # Log non-default camera targets so job history shows which view ran. - default_display = multicam_params.get('multicam_default_display') - if not default_display or camera_name != default_display: - print(f'Running pipeline on camera: {camera_name}') - with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): - _working_directory_path = Path(_working_directory) - input_path = utils.make_directory(_working_directory_path / 'input') - trained_pipeline_path = utils.make_directory(_working_directory_path / 'trained_pipeline') - output_path = utils.make_directory(_working_directory_path / 'output') - - detector_output_file = str(output_path / 'detector_output.csv') - track_output_file = str(output_path / 'track_output.csv') - img_list_path = input_path / 'img_list_file.txt' - - pipeline_path = _resolve_pipeline_path(conf, gc, pipeline, trained_pipeline_path) - - assert pipeline_path.exists(), ( - "Requested pipeline could not be found." - " Make sure that VIAME is installed correctly and all addons have loaded." - f" Job asked for {pipeline_path} but it does not exist" - ) - - if multicam_cameras: - input_folder = gc.getFolder(input_folder_id) - input_fps = fromMeta(input_folder, constants.FPSMarker) - requires_input = multicam_params.get('multicam_requires_input', False) - creates_new_dataset = pipeline_creates_new_dataset(pipeline) - camera_media: Dict[str, Tuple[List[str], str]] = {} - - for cam_index, camera in enumerate(multicam_cameras, start=1): - cam_input_path = utils.make_directory(input_path / camera['name']) - media_list, media_type = utils.download_source_media( - gc, camera['folder_id'], cam_input_path, force_transcoded - ) - if frame_range is not None and media_type == constants.ImageSequenceType: - media_list = filter_image_list_by_frame_range(media_list, frame_range) - camera_media[camera['name']] = (media_list, media_type) - if requires_input and camera.get('input_revision') is not None: - gt_path = _working_directory_path / f'detections{cam_index}.csv' - utils.download_revision_csv( - gc, camera['folder_id'], camera['input_revision'], gt_path - ) - - arg_file_pair, out_files = build_multicam_kwiver_settings( - _working_directory_path, - multicam_cameras, - camera_media, - requires_input=requires_input, - ) - - command = [ - f". {shlex.quote(str(conf.viame_setup_script))} &&", - f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", - "viame runner", - f"-p {shlex.quote(str(pipeline_path))}", - ] - if input_type == constants.VideoType: - command.extend( - [ - '-s input:video_reader:type=vidl_ffmpeg', - f"-s downsampler:target_frame_rate={shlex.quote(str(input_fps))}", - ] - ) - if frame_range is not None: - _append_frame_range_video_settings( - command, input_folder, frame_range, pipeline['pipe'] - ) - for arg, file_name in arg_file_pair.items(): - command.append(f"-s {shlex.quote(arg)}={shlex.quote(file_name)}") - - transcoded_video: Optional[str] = None - if creates_new_dataset: - video_name = None - if is_transcode_pipeline(pipeline): - video_name = str( - output_path / f"{pipeline.get('name', 'transcode')}_{input_folder_id}.mp4" - ) - transcoded_video = append_new_dataset_media_writers( - command, pipeline, output_path, video_filename=video_name - ) - - calibration_item_id = multicam_params.get('calibration_item_id') - if calibration_item_id and is_stereo_measurement_pipeline(pipeline): - cal_item = gc.getItem(calibration_item_id) - cal_dir = utils.make_directory(_working_directory_path / 'calibration') - gc.downloadItem( - calibration_item_id, - str(cal_dir), - name=cal_item.get('name'), - ) - cal_path = find_downloaded_calibration_file(cal_dir) - if cal_path is not None: - append_stereo_calibration_kwiver_settings(command, cal_path, pipeline) - else: - manager.write( - f'Warning: calibration item {calibration_item_id} ' - f'has no recognized calibration file under {cal_dir}\n' - ) - - # One image list per camera (each a single line-separated file). - input_manifests = [ - arg_file_pair[f'input{i + 1}:video_filename'] - for i in range(len(multicam_cameras)) - if f'input{i + 1}:video_filename' in arg_file_pair - ] - _append_input_list_kwiver_settings(command, pipeline, input_manifests) - - _inject_dataset_metadata_file(command, gc, _working_directory_path, params, manager) - - kwiver_params = params.get('kwiver_params') - if kwiver_params: - for key, value in kwiver_params.items(): - command.append(f'-s {shlex.quote(key)}={shlex.quote(str(value))}') - - manager.updateStatus(JobStatus.RUNNING) - popen_kwargs = { - 'args': " ".join(command), - 'shell': True, - 'executable': '/bin/bash', - 'cwd': output_path, - 'env': conf.gpu_process_env, - } - utils.stream_subprocess(self, context, manager, popen_kwargs) - - if ( - is_stereo_measurement_pipeline(pipeline) - and 'calibrate_cameras' in str(pipeline.get('pipe', '')).lower() - ): - calibration_outputs = _find_stereo_calibration_outputs(output_path) - if calibration_outputs: - calibration_output = calibration_outputs[0] - try: - uploaded_calibration = gc.uploadFileToFolder( - input_folder_id, - str(calibration_output), - ) - uploaded_calibration_file_id = uploaded_calibration.get('_id') - if uploaded_calibration_file_id is not None: - uploaded_calibration_file_id_str = str(uploaded_calibration_file_id) - cal_url = ( - f'/dive_dataset/{input_folder_id}/calibration' - f'?fileId={uploaded_calibration_file_id_str}' - ) - gc.sendRestRequest('POST', cal_url) - manager.write( - 'Assigned calibration output to dataset: ' - f'{calibration_output.name}\n' - ) - else: - manager.write( - 'Warning: uploaded calibration output ' - f'{calibration_output.name} has no file id\n' - ) - except Exception as exc: - manager.write( - 'Warning: failed to assign calibration output ' - f'{calibration_output.name}: {exc}\n' - ) - else: - manager.write( - 'Warning: stereo calibration pipeline produced no ' - 'recognized calibration output file\n' - ) - - if creates_new_dataset: - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - _push_new_dataset_from_media( - gc, - manager, - params, - input_folder_id, - output_path, - pipeline, - transcoded_video=transcoded_video, - ) - return - - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - for camera in multicam_cameras: - cam_name = camera['name'] - output_name = out_files[cam_name] - # Multicam KWIVER args use basename-only writers; viame cwd is output_path, - # so CSVs are created under output/, not the temp directory root. - output_file = output_path / output_name - if not output_file.exists() or not output_file.stat().st_size: - detector_name = output_name.replace('computed_tracks', 'computed_detections') - detector_path = output_path / detector_name - if detector_path.exists() and detector_path.stat().st_size: - output_file = detector_path - if frame_range is not None and camera_media[cam_name][1] == constants.VideoType: - filtered_path = filter_csv_by_frame_range(str(output_file), frame_range) - output_file = Path(filtered_path) - newfile = gc.uploadFileToFolder(camera['folder_id'], str(output_file)) - gc.addMetadataToItem(str(newfile["itemId"]), {"pipeline": pipeline}) - gc.post( - f'dive_rpc/postprocess/{camera["folder_id"]}', - data={"skipJobs": True}, - ) - return - - # Download source media - input_folder: GirderModel = gc.getFolder(input_folder_id) - creates_new_dataset = pipeline_creates_new_dataset(pipeline) - input_media_list, _ = utils.download_source_media( - gc, input_folder_id, input_path, force_transcoded - ) - - if input_type == constants.VideoType: - input_fps = fromMeta(input_folder, constants.FPSMarker) - assert len(input_media_list) == 1, "Expected exactly 1 video" - command = [ - f". {shlex.quote(str(conf.viame_setup_script))} &&", - f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", - "viame runner", - "-s input:video_reader:type=vidl_ffmpeg", - f"-p {shlex.quote(str(pipeline_path))}", - f"-s input:video_filename={shlex.quote(input_media_list[0])}", - f"-s downsampler:target_frame_rate={shlex.quote(str(input_fps))}", - f"-s detector_writer:file_name={shlex.quote(detector_output_file)}", - f"-s track_writer:file_name={shlex.quote(track_output_file)}", - ] - if frame_range is not None: - _append_frame_range_video_settings( - command, input_folder, frame_range, pipeline['pipe'] - ) - elif input_type == constants.ImageSequenceType: - # Filter image list by frame range if specified - filtered_media_list = input_media_list - if frame_range is not None: - filtered_media_list = filter_image_list_by_frame_range( - input_media_list, frame_range - ) - with open(img_list_path, "w+") as img_list_file: - img_list_file.write('\n'.join(filtered_media_list)) - command = [ - f". {shlex.quote(str(conf.viame_setup_script))} &&", - f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", - "viame runner", - f"-p {shlex.quote(str(pipeline_path))}", - f"-s input:video_filename={shlex.quote(str(img_list_path))}", - f"-s detector_writer:file_name={shlex.quote(detector_output_file)}", - f"-s track_writer:file_name={shlex.quote(track_output_file)}", - ] - else: - raise ValueError('Unknown input type: {}'.format(input_type)) - - # Include input detections - if input_revision is not None: - pipeline_input_file = input_path / 'groundtruth.csv' - utils.download_revision_csv(gc, input_folder_id, input_revision, pipeline_input_file) - quoted_input_file = shlex.quote(str(pipeline_input_file)) - command.append(f'-s detection_reader:file_name={quoted_input_file}') - command.append(f'-s track_reader:file_name={quoted_input_file}') - - transcoded_video = None - if creates_new_dataset: - video_name = None - if is_transcode_pipeline(pipeline): - video_name = str( - output_path / f"{pipeline.get('name', 'transcode')}_{input_folder_id}.mp4" - ) - transcoded_video = append_new_dataset_media_writers( - command, pipeline, output_path, video_filename=video_name - ) - - single_input_manifest = ( - str(img_list_path) if input_type == constants.ImageSequenceType else input_media_list[0] - ) - _append_input_list_kwiver_settings(command, pipeline, [single_input_manifest]) - - _inject_dataset_metadata_file(command, gc, _working_directory_path, params, manager) - - # Apply user-provided KWIVER parameter overrides. - kwiver_params = params.get('kwiver_params') - if kwiver_params: - for key, value in kwiver_params.items(): - command.append(f'-s {shlex.quote(key)}={shlex.quote(str(value))}') - - manager.updateStatus(JobStatus.RUNNING) - popen_kwargs = { - 'args': " ".join(command), - 'shell': True, - 'executable': '/bin/bash', - 'cwd': output_path, - 'env': conf.gpu_process_env, - } - utils.stream_subprocess(self, context, manager, popen_kwargs) - - if creates_new_dataset: - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - _push_new_dataset_from_media( - gc, - manager, - params, - input_folder_id, - output_path, - pipeline, - transcoded_video=transcoded_video, - ) - return - - if Path(track_output_file).exists() and os.path.getsize(track_output_file): - output_file = track_output_file - else: - output_file = detector_output_file - - # Filter output CSV by frame range for videos - if frame_range is not None and input_type == constants.VideoType: - output_file = filter_csv_by_frame_range(output_file, frame_range) - - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - newfile = gc.uploadFileToFolder(output_folder_id, output_file) - - gc.addMetadataToItem(str(newfile["itemId"]), {"pipeline": pipeline}) - gc.post(f'dive_rpc/postprocess/{output_folder_id}', data={"skipJobs": True}) - - -@app.task(bind=True, acks_late=True, ignore_results=True) -def export_trained_pipeline(self: Task, params: ExportTrainedPipelineJob): - conf = Config() - conf.require_viame_install() - context: dict = {} - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - gc: GirderClient = self.girder_client - utils.authenticate_urllib(gc) - manager.updateStatus(JobStatus.FETCHING_INPUT) - - # Extract params - input_folder_id = params["input_folder"] - output_folder_id = params["output_folder"] - output_name = params["output_name"] - - with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): - _working_directory_path = Path(_working_directory) - trained_pipeline_path = utils.make_directory(_working_directory_path / 'trained_pipeline') - output_path = utils.make_directory(_working_directory_path / 'output') - onnx_path = output_path / output_name - convert_to_onnx_pipeline_path = conf.viame_pipeline_path / "convert_model_to_onnx.pipe" - - gc.downloadFolderRecursive(input_folder_id, str(trained_pipeline_path)) - extensions = ['*.weights', '*.ckpt', '*.pth'] - model_file = None - - for ext in extensions: - found_files = list(trained_pipeline_path.glob(ext)) - if found_files: - model_file = found_files[0] - break - - if not model_file: - raise FileNotFoundError(f"No weights path ({extensions}) found.") - - # Convert pipeline to ONNX - command = [ - f". {shlex.quote(str(conf.viame_setup_script))} &&", - f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", - "viame runner", - f"-p {shlex.quote(str(convert_to_onnx_pipeline_path))}", - f"-s onnx_convert:model_path={shlex.quote(str(model_file))}", - f"-s onnx_convert:onnx_model_prefix={shlex.quote(str(onnx_path))}", - ] - - manager.updateStatus(JobStatus.RUNNING) - popen_kwargs = { - 'args': " ".join(command), - 'shell': True, - 'executable': '/bin/bash', - 'cwd': output_path, - 'env': conf.gpu_process_env, - } - utils.stream_subprocess(self, context, manager, popen_kwargs) - - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - gc.uploadFileToFolder(output_folder_id, onnx_path) - - -@app.task(bind=True, acks_late=True, ignore_result=True) -def train_pipeline(self: Task, params: TrainingJob): - """Train a pipeline by making a call to viame train""" - conf = Config() - conf.require_viame_install() - context: dict = {} - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - gc: GirderClient = self.girder_client - utils.authenticate_urllib(gc) - manager.updateStatus(JobStatus.FETCHING_INPUT) - - # Extract params - results_folder_id = params['results_folder_id'] - dataset_input_list = params['dataset_input_list'] - pipeline_name = params['pipeline_name'] - config = params['config'] - annotated_frames_only = params['annotated_frames_only'] - label_text = params['label_txt'] - model = params.get('model', None) - # Normalize: model can arrive as a list of [key, value] pairs from some serialization paths - if model is not None and isinstance(model, list): - model = dict(model) - force_transcoded = params.get('force_transcoded', False) - - pipeline_base_path = Path(conf.get_extracted_pipeline_path()) - config_file = pipeline_base_path / config - # List of (input folder, ground truth file) pairs for creating input lists - input_groundtruth_list: List[Tuple[Path, Path]] = [] - # root_data_dir is the directory passed to `viame train` - with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): - _working_directory_path = Path(_working_directory) - input_path = utils.make_directory(_working_directory_path / 'input') - output_path = utils.make_directory(_working_directory_path / 'output') - - for source_folder_id, revision in dataset_input_list: - download_path = utils.make_directory(input_path / source_folder_id) - groundtruth_path = download_path / 'groundtruth.csv' - # Download groundtruth item - utils.download_revision_csv(gc, source_folder_id, revision, groundtruth_path) - # Download input media - input_media_list, input_type = utils.download_source_media( - gc, source_folder_id, download_path, force_transcoded - ) - if input_type == constants.VideoType: - download_path = Path(input_media_list[0]) - # Set media source location - input_groundtruth_list.append((download_path, groundtruth_path)) - - input_folder_file_list = input_path / "input_folder_list.txt" - ground_truth_file_list = input_path / "input_truth_list.txt" - with open(input_folder_file_list, "w+") as data_list: - with open(ground_truth_file_list, "w+") as truth_list: - for folder_path, groundtruth_path in input_groundtruth_list: - data_list.write(f"{folder_path}\n") - truth_list.write(f"{groundtruth_path}\n") - - training_results_path = utils.make_directory(output_path / "category_models") - - command = [ - f". {shlex.quote(str(conf.viame_setup_script))} &&", - f"KWIVER_DEFAULT_LOG_LEVEL={shlex.quote(conf.kwiver_log_level)}", - f"{shlex.quote(str(conf.viame_executable))} train", - "--input-list", - shlex.quote(str(input_folder_file_list)), - "--input-truth", - shlex.quote(str(ground_truth_file_list)), - "--config", - shlex.quote(str(config_file)), - "--no-query", - "--no-embedded-pipe", - ] - - if annotated_frames_only: - command.append("--gt-frames-only") - - if label_text: - labels_path = input_path / "labels.txt" - with open(labels_path, "w+") as labels_file: - labels_file.write(label_text) - command.append("--labels") - command.append(shlex.quote(str(labels_path))) - - if model: - model_path = None - if model.get('folderId', False): - trained_pipeline_path = utils.make_directory( - _working_directory_path / 'trained_pipeline' - ) - gc.downloadFolderRecursive(model["folderId"], str(trained_pipeline_path)) - model_path = trained_pipeline_path / model["name"] - elif model.get('path', False): - model_path = model['path'] - if model_path: - command.append("--init-weights") - command.append(shlex.quote(str(model_path))) - - manager.updateStatus(JobStatus.RUNNING) - popen_kwargs = { - 'args': " ".join(command), - 'shell': True, - 'executable': '/bin/bash', - 'cwd': output_path, - 'env': conf.gpu_process_env, - } - utils.stream_subprocess(self, context, manager, popen_kwargs) - - # Check that there are results in the output path - if len(list(training_results_path.glob("*"))) == 0: - raise RuntimeError("Training output didn't produce results, discarding...") - - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - # This is the name of the folder that is uploaded to the - # "Training Results" girder folder - girder_output_folder = gc.createFolder( - results_folder_id, - pipeline_name, - metadata={ - constants.TrainedPipelineMarker: True, - "trained_on": dataset_input_list, - }, - ) - gc.upload(f"{training_results_path}/*", girder_output_folder["_id"]) - - -@app.task(bind=True, acks_late=True, ignore_result=True) -def convert_calibration(self: Task, itemId: str): - """ - Convert a calibrationFile item to a JSON camera-rig in a separate Girder item - marked jsonCalibrationFile for display. - """ - conf = Config() - conf.require_viame_install() - context: dict = {} - gc: GirderClient = self.girder_client - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - convert_tool = conf.viame_install_path / 'configs' / 'convert_cam_format.py' - - with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): - _working_directory_path = Path(_working_directory) - item: GirderModel = gc.getItem(itemId) - folder_id = str(item.get('folderId')) - folder = gc.getFolder(folder_id) - multi_cam = (folder.get('meta') or {}).get(constants.MultiCamMarker) or {} - if str(multi_cam.get(constants.CalibrationItemIdMarker)) != str(itemId): - manager.write('Calibration source was replaced; skipping stale conversion job.\n') - return - - files = list(gc.listFile(itemId)) - source_file = next( - (f for f in files if constants.stereoCalibrationRegex.search(f['name'])), - None, - ) - if source_file is None: - manager.write('No convertible calibration file found in item; skipping.\n') - return - - manager.updateStatus(JobStatus.FETCHING_INPUT) - gc.downloadItem(itemId, _working_directory_path, name=item.get('name')) - input_path = _working_directory_path / source_file['name'] - file_bytes = input_path.read_bytes() - if source_file['name'].lower().endswith('.json'): - if calibration_format.calibration_upload_is_final_json(source_file['name'], file_bytes): - manager.write('Source calibration is already JSON; skipping conversion.\n') - return - input_path = calibration_format.prepare_conversion_input_path( - input_path, - file_bytes, - _working_directory_path, - ) - output_path = input_path.with_suffix('.json') - - manager.updateStatus(JobStatus.RUNNING) - command = [ - f". {shlex.quote(str(conf.viame_setup_script))} &&", - "python", - shlex.quote(str(convert_tool)), - shlex.quote(str(input_path)), - shlex.quote(str(output_path)), - ] - popen_kwargs = { - 'args': " ".join(command), - 'shell': True, - 'executable': '/bin/bash', - 'cwd': str(_working_directory_path), - 'env': conf.gpu_process_env, - } - try: - utils.stream_subprocess(self, context, manager, popen_kwargs) - except Exception as exc: - error_msg = str(exc) or 'Calibration conversion failed' - updated_multi_cam = dict(multi_cam) - updated_multi_cam[constants.CalibrationConversionErrorMarker] = error_msg - gc.addMetadataToFolder( - folder_id, - {constants.MultiCamMarker: updated_multi_cam}, - ) - raise - - if not output_path.exists() or not output_path.stat().st_size: - error_msg = 'Calibration conversion produced no JSON output' - updated_multi_cam = dict(multi_cam) - updated_multi_cam[constants.CalibrationConversionErrorMarker] = error_msg - gc.addMetadataToFolder( - folder_id, - {constants.MultiCamMarker: updated_multi_cam}, - ) - raise RuntimeError(error_msg) - - folder = gc.getFolder(folder_id) - multi_cam = (folder.get('meta') or {}).get(constants.MultiCamMarker) or {} - if str(multi_cam.get(constants.CalibrationItemIdMarker)) != str(itemId): - manager.write('Calibration source was replaced during conversion; discarding output.\n') - return - - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - json_name = output_path.name - gc.upload(str(output_path), folder_id) - json_items = sorted( - gc.listItem(folder_id, name=json_name), - key=lambda existing: existing.get('created', ''), - ) - if not json_items: - raise RuntimeError('Failed to create calibration JSON item') - json_item = json_items[-1] - json_item_id = str(json_item['_id']) - gc.addMetadataToItem( - json_item_id, - { - constants.JsonCalibrationFileMarker: 'true', - constants.CalibrationFileMarker: False, - }, - ) - - folder = gc.getFolder(folder_id) - multi_cam = (folder.get('meta') or {}).get(constants.MultiCamMarker) or {} - if str(multi_cam.get(constants.CalibrationItemIdMarker)) != str(itemId): - manager.write( - 'Calibration source was replaced before linking JSON; discarding output.\n' - ) - gc.delete(f'item/{json_item_id}') - return - - updated_multi_cam = dict(multi_cam) - updated_multi_cam.setdefault(constants.CalibrationItemIdMarker, str(itemId)) - updated_multi_cam[constants.JsonCalibrationItemIdMarker] = json_item_id - updated_multi_cam.setdefault( - constants.CalibrationOriginalNameMarker, - source_file['name'], - ) - updated_multi_cam.pop(constants.CalibrationConversionErrorMarker, None) - gc.addMetadataToFolder( - folder_id, - {constants.MultiCamMarker: updated_multi_cam}, - ) - - for existing_item in gc.listItem(folder_id): - existing_id = str(existing_item['_id']) - if existing_id in {json_item_id, str(itemId)}: - continue - existing_meta = existing_item.get('meta') or {} - if asbool(existing_meta.get(constants.JsonCalibrationFileMarker)): - gc.delete(f"item/{existing_id}") - - -def resolve_annotation_fps( - gc: GirderClient, - folder_id: str, - *, - native_fps: Optional[float] = None, - default_fps: float = 1.0, -) -> float: - """Pick annotation FPS from current folder meta vs media FPS. - - Re-reads folder ``fps`` so a concurrent CSV import (assetstore postprocess) - is not overwritten by a stale ``-1`` snapshot from job start. - - For video, pass ``native_fps`` from ffprobe. For image sequences, omit it so - ``-1`` falls back to ``default_fps`` (1) and any CSV-set value is kept. - """ - requested_fps = fromMeta(gc.getFolder(folder_id), constants.FPSMarker) - return utils.choose_annotation_fps( - requested_fps, native_fps=native_fps, default_fps=default_fps - ) - - -def _download_video_item( - gc: GirderClient, - manager: JobManager, - item_id: str, - item_name: str, - dest_dir: Path, -) -> str: - """Download a Girder video item to *dest_dir*; return the local file path.""" - file_name = str(dest_dir / item_name) - manager.updateStatus(JobStatus.FETCHING_INPUT) - manager.write(f'Fetching input from {item_id} to {file_name}...\n') - gc.downloadItem(item_id, dest_dir, name=item_name) - return file_name - - -@app.task(bind=True, acks_late=True, ignore_result=True) -def convert_video( - self: Task, folderId: str, itemId: str, user_id: str, user_login: str, skip_transcoding=False -): - context: dict = {} - gc: GirderClient = self.girder_client - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): - _working_directory_path = Path(_working_directory) - item: GirderModel = gc.getItem(itemId) - item_name = item['name'] - output_file_path = (_working_directory_path / item_name).with_suffix('.transcoded.mp4') - - # When skip_transcoding is requested, probe via authenticated HTTP Range - # requests first so web-ready S3/filesystem videos never need a full download. - # Fall back to downloading the whole object if remote probe/alignment fails. - jsoninfo = None - file_name: Optional[str] = None - auth_headers: Optional[str] = None - remote_url: Optional[str] = None - - if skip_transcoding: - try: - remote_url = utils.item_primary_file_download_url(gc, itemId) - except Exception as exc: - manager.write( - f'Could not resolve file download URL ({exc}); ' - 'falling back to full download\n' - ) - remote_url = None - if remote_url is not None: - auth_headers = utils.girder_auth_headers(gc.token) - manager.updateStatus(JobStatus.RUNNING) - manager.write(f'Probing video via HTTP Range requests: {remote_url}\n') - try: - jsoninfo = utils.ffprobe_format_and_streams( - self, context, manager, remote_url, headers=auth_headers - ) - except utils.CanceledError: - raise - except Exception as exc: - manager.write( - f'Remote ffprobe failed ({exc}); falling back to full download\n' - ) - jsoninfo = None - auth_headers = None - remote_url = None - - if jsoninfo is None: - file_name = _download_video_item( - gc, manager, itemId, item_name, _working_directory_path - ) - manager.updateStatus(JobStatus.RUNNING) - jsoninfo = utils.ffprobe_format_and_streams(self, context, manager, file_name) - - videostream = list(filter(lambda x: x["codec_type"] == "video", jsoninfo["streams"])) - if len(videostream) != 1: - print('Expected 1 video stream, found {}'.format(len(videostream))) - print('Using first Video Stream found') - - format_info = jsoninfo.get('format') or {} - format_name = format_info.get('format_name') or '' - - # Extract framerate (avg_frame_rate, else r_frame_rate for e.g. MPEG-TS) - originalFpsString, originalFps = utils.fps_from_ffprobe_stream(videostream[0]) - - source_misaligned = False - if skip_transcoding: - alignment_source: Optional[str] = file_name or remote_url - alignment_headers = auth_headers if file_name is None else None - try: - source_misaligned = is_frame_misaligned( - self, - alignment_source, - context, - manager, - headers=alignment_headers, - ) - except utils.CanceledError: - raise - except Exception as exc: - if file_name is None: - manager.write( - f'Remote frame-alignment check failed ({exc}); ' - 'falling back to full download\n' - ) - file_name = _download_video_item( - gc, manager, itemId, item_name, _working_directory_path - ) - manager.updateStatus(JobStatus.RUNNING) - # Re-probe locally so metadata matches the file we will encode. - jsoninfo = utils.ffprobe_format_and_streams(self, context, manager, file_name) - videostream = list( - filter(lambda x: x["codec_type"] == "video", jsoninfo["streams"]) - ) - format_info = jsoninfo.get('format') or {} - format_name = format_info.get('format_name') or '' - originalFpsString, originalFps = utils.fps_from_ffprobe_stream(videostream[0]) - source_misaligned = is_frame_misaligned(self, Path(file_name), context, manager) - else: - raise - - # Skip remux/transcode only for browser-safe sources, matching desktop checks. - can_skip_transcode = utils.can_skip_video_transcoding( - skip_transcoding=skip_transcoding, - codec_name=videostream[0]['codec_name'], - sample_aspect_ratio=videostream[0].get('sample_aspect_ratio'), - format_name=format_name, - source_misaligned=source_misaligned, - ) - - # lets determine if we don't need to transcode this file - if can_skip_transcode: - # Now we can update the meta data and push the values - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - if file_name is None: - manager.write('Skip transcode: no full download required\n') - newAnnotationFps = resolve_annotation_fps(gc, folderId, native_fps=originalFps) - gc.addMetadataToItem( - itemId, - { - "source_video": False, # even though it is, this for requesting - "transcoder": "ffmpeg", - constants.OriginalFPSMarker: originalFps, - constants.OriginalFPSStringMarker: originalFpsString, - "codec": "h264", - }, - ) - gc.addMetadataToFolder( - folderId, - { - constants.DatasetMarker: True, # mark the parent folder as able to annotate. - constants.OriginalFPSMarker: originalFps, - constants.OriginalFPSStringMarker: originalFpsString, - constants.FPSMarker: newAnnotationFps, - "ffprobe_info": videostream[0], - }, - ) - return - elif skip_transcoding: - print('Transcoding cannot be skipped:') - print(f'Codec Name: {videostream[0]["codec_name"]}') - print(f'format_name: {format_name}') - if videostream[0]['codec_name'] != 'h264': - print('Codec is not h264; file will be transcoded') - elif videostream[0].get('sample_aspect_ratio') != '1:1': - print( - 'Sample aspect ratio is not 1:1; file will be transcoded ' - '(desktop-parity rule)' - ) - elif not utils.container_allows_skip_transcoding(format_name): - print('Container is not web-safe (e.g. mpegts); file will be transcoded') - elif source_misaligned: - print('Frame timestamps are misaligned; file will be transcoded') - - if file_name is None: - file_name = _download_video_item( - gc, manager, itemId, item_name, _working_directory_path - ) - manager.updateStatus(JobStatus.RUNNING) - - command = [ - "ffmpeg", - "-i", - file_name, - "-c:v", - "libx264", - "-preset", - "slow", - # https://github.com/Kitware/dive/issues/855 - "-crf", - "22", - # https://askubuntu.com/questions/1315697/could-not-find-tag-for-codec-pcm-s16le-in-stream-1-codec-not-currently-support - "-c:a", - "aac", - # see native/ code for a discussion of this option - "-vf", - "scale=ceil(iw*sar/2)*2:ceil(ih/2)*2,setsar=1", - str(output_file_path), - ] - utils.stream_subprocess(self, context, manager, {'args': command}) - # Check to see if frame alignment remains the same - aligned_file = check_and_fix_frame_alignment(self, output_file_path, context, manager) - misaligned_flag = False - if aligned_file != output_file_path: - misaligned_flag = True - - manager.updateStatus(JobStatus.PUSHING_OUTPUT) - newAnnotationFps = resolve_annotation_fps(gc, folderId, native_fps=originalFps) - new_file = gc.uploadFileToFolder(folderId, aligned_file) - gc.addMetadataToItem( - new_file['itemId'], - { - "source_video": False, - "transcoder": "ffmpeg", - constants.OriginalFPSMarker: originalFps, - constants.OriginalFPSStringMarker: originalFpsString, - "codec": "h264", - }, - ) - source_metadata = { - "source_video": True, - constants.OriginalFPSMarker: originalFps, - constants.OriginalFPSStringMarker: originalFpsString, - "codec": videostream[0]["codec_name"], - } - if misaligned_flag: - source_metadata[constants.MISALGINED_MARKER] = True - gc.addMetadataToItem( - itemId, - source_metadata, - ) - gc.addMetadataToFolder( - folderId, - { - constants.DatasetMarker: True, # mark the parent folder as able to annotate. - constants.OriginalFPSMarker: originalFps, - constants.OriginalFPSStringMarker: originalFpsString, - constants.FPSMarker: newAnnotationFps, - "ffprobe_info": videostream[0], - }, - ) - - -@app.task(bind=True, acks_late=True) -def convert_images(self: Task, folderId, user_id: str, user_login: str): - """ - Ensures that all images in a folder are in a web friendly format (png or jpeg). - - If conversions succeeds for an image, it will replace the image with an image - of the same name, but in a web friendly extension. - - Returns the number of images successfully converted. - """ - context: dict = {} - gc: GirderClient = self.girder_client - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - items_to_convert = [ - item - for item in gc.listItem(folderId) - if ( - constants.imageRegex.search(item["name"]) - and not constants.safeImageRegex.search(item["name"]) - ) - ] - - with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): - working_directory_path = Path(_working_directory) - images_path = utils.make_directory(working_directory_path / 'images') - - for item in items_to_convert: - # Assumes 1 file per item - gc.downloadItem(item["_id"], images_path, item["name"]) - - item_path = images_path / item["name"] - new_item_path = images_path / ".".join([*item["name"].split(".")[:-1], "png"]) - command = ["ffmpeg", "-i", str(item_path), str(new_item_path)] - utils.stream_subprocess(self, context, manager, {'args': command}) - gc.uploadFileToFolder(folderId, new_item_path) - gc.delete(f"item/{str(item['_id'])}") - - gc.addMetadataToFolder( - str(folderId), - { - "annotate": True, # mark the parent folder as able to annotate. - constants.FPSMarker: resolve_annotation_fps(gc, folderId), - }, - ) - - -@app.task(bind=True, acks_late=True) -def convert_large_images(self: Task, folderId, user_id: str, user_login: str): - """ - Converts all images in the folder to large images - - This is typically done if the images are >8k W or L resolution - - Returns the number of images successfully converted. - """ - context: dict = {} - gc: GirderClient = self.girder_client - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - items_to_convert = [ - item for item in gc.listItem(folderId) if (constants.safeImageRegex.search(item["name"])) - ] - for item in items_to_convert: - # Assumes 1 file per item - try: - # Does it already have tiles? - gc.get(f'item/{item["_id"]}/tiles') - manager.write(f'Skipping {item["name"]}, already a large image\n') - continue - except HttpError as e: - # Safely parse JSON if possible - message = "" - try: - message = e.response.json().get("message", "") - except Exception: - pass # non-JSON response, leave message empty - # This is the Girder message when no large image exists - if e.status == 400 and message == "No large image file in this item.": - manager.write(f'Converting {item["name"]} to large image\n') - gc.post(f'item/{item["_id"]}/tiles') - else: - # Re-raise unexpected errors to fail the job - raise - gc.addMetadataToFolder( - str(folderId), - {"type": constants.LargeImageType}, # mark the parent folder as able to annotate. - ) - - -@app.task(bind=True, acks_late=True, ignore_result=True) -def extract_zip(self: Task, folderId: str, itemId: str, user_id: str, user_login: str): - """ - Discovery logic: - * Find all folders that have at least one child file (potential datasets) - * Exclude folders which are sub-folders of previously discovered folders - because datasets cannot be nested in other datasets - """ - context: dict = {} - gc: GirderClient = self.girder_client - manager: JobManager = patch_manager(self.job_manager) - if utils.check_canceled(self, context): - manager.updateStatus(JobStatus.CANCELED) - return - - with tempfile.TemporaryDirectory() as _working_directory, suppress(utils.CanceledError): - _working_directory_path = Path(_working_directory) - item: GirderModel = gc.getItem(itemId) - file_name = str(_working_directory_path / item['name']) - manager.write(f'Fetching input from {itemId} to {file_name}...\n') - gc.downloadItem(itemId, _working_directory, item["name"]) - discovered_folders = {} - with zipfile.ZipFile(file_name, 'r') as zipObj: - listOfFileNames = zipObj.namelist() - sum_file_size = sum([data.file_size for data in zipObj.filelist]) - sum_compress_size = sum([data.compress_size for data in zipObj.filelist]) - ratio = sum_file_size / sum_compress_size - if ratio > 600: - manager.write(f"Compression ratio is exceedingly high at {ratio}\n\ - Please contact an admin at viame-web@kitware.com if this is a valid zip file") - raise Exception("High Compression Ratio for Zip File") - - multicam_export_roots = { - os.path.dirname(fileName) - for fileName in listOfFileNames - if os.path.basename(fileName) == constants.MultiCamJsonFileName - and not fileName.endswith(os.path.sep) - } - - for fileName in listOfFileNames: - folderName = os.path.dirname(fileName) - parentName = os.path.dirname(folderName) - if parentName in discovered_folders and folderName != '': - discovered_folders[folderName] = 'ignored' - # Nested single-camera exports stay skipped; multicam camera trees must extract. - if not utils.is_path_under_multicam_export(folderName, multicam_export_roots): - continue - if fileName.endswith(os.path.sep): - continue - if folderName not in discovered_folders: - discovered_folders[folderName] = 'unstructured' - if constants.metaRegex.search(os.path.basename(fileName)): - if folderName in multicam_export_roots: - discovered_folders[folderName] = 'multicam' - else: - discovered_folders[folderName] = 'dataset' - if fileName.endswith('.zip'): - raise Exception("Nested Zip Files are invalid") - manager.write(f"Extracting: {fileName}\n") - zipObj.extract(fileName, f'{_working_directory}') - - # remove the zip file so it isn't uploaded back to the folder - os.remove(file_name) - # Create source folder and move zip file there - created_folder = gc.createFolder( - folderId, - constants.SourceFolderName, - reuseExisting=True, - ) - gc.sendRestRequest( - "PUT", - f"/item/{str(item['_id'])}?folderId={str(created_folder['_id'])}", - ) - # Only make subfolders if more than 1 discovered folder exists - make_subfolders = ( - len(discovered_folders) - list(discovered_folders.values()).count('ignored') - ) > 1 - for folderName, folderType in discovered_folders.items(): - subFolderName = folderName if make_subfolders else '' - if folderType == 'unstructured': - utils.upload_zipped_flat_media_files( - gc, - manager, - folderId, - _working_directory_path / folderName, - subFolderName, - ) - elif folderType == 'multicam': - utils.upload_exported_multicam_zipped_dataset( - gc, - manager, - folderId, - _working_directory_path / folderName, - subFolderName, - ) - elif folderType == 'dataset': - utils.upload_exported_zipped_dataset( - gc, - manager, - folderId, - _working_directory_path / folderName, - subFolderName, - ) - else: - manager.write(f'Ignoring {folderName}\n') - - if make_subfolders: - gc.sendRestRequest( - "DELETE", - f"folder/{folderId}/metadata", - json=[constants.TypeMarker, constants.FPSMarker, constants.DatasetMarker], - ) diff --git a/server/dive_tasks/upgrade_pipelines.py b/server/dive_tasks/upgrade_pipelines.py new file mode 100644 index 000000000..5be5b1cfa --- /dev/null +++ b/server/dive_tasks/upgrade_pipelines.py @@ -0,0 +1,140 @@ +import logging +import os +from pathlib import Path +import shutil +from typing import List +from urllib import request +from urllib.parse import urlparse +import zipfile + +import gdown +from gdown.parse_url import is_google_drive_url, parse_url +from girder_client import GirderClient +from girder_worker.app import app +from girder_worker.task import Task +from girder_worker.utils import JobManager, JobStatus + +from dive_tasks import utils +from dive_tasks.manager import patch_manager +from dive_tasks.pipeline_discovery import discover_configs +from dive_tasks.viame_config import EMPTY_JOB_SCHEMA, Config + +logger = logging.getLogger(__name__) + +# https://github.com/VIAME/VIAME/blob/master/cmake/download_viame_addons.csv +UPGRADE_JOB_DEFAULT_URLS: List[str] = [ + 'https://viame.kitware.com/api/v1/item/627b145487bad2e19a4c4697/download', # HabCam + 'https://viame.kitware.com/api/v1/item/627b32b1994809b024f207a7/download', # SEFSC + 'https://viame.kitware.com/api/v1/item/627b3289ea630db5587b577d/download', # SWFSC-PengHead + 'https://viame.kitware.com/api/v1/item/627b326fea630db5587b577b/download', # Motion + 'https://viame.kitware.com/api/v1/item/627b326cc4da86e2cd3abb5b/download', # EM Tuna + 'https://viame.kitware.com/api/v1/item/627b3282c4da86e2cd3abb5d/download', # MOUSS + 'https://viame.kitware.com/api/v1/item/615bc7aa7e5c13a5bb9af7a7/download', # Aerial Penguin + 'https://viame.kitware.com/api/v1/item/629807c192adc2f0ecfa5b54/download', # Sea Lion +] + + +def _normalize_google_drive_url(url: str) -> str: + """Strip a leading www. so gdown recognizes common pasted Drive links.""" + parsed = urlparse(url) + host = parsed.netloc.lower() + if host.startswith('www.'): + return parsed._replace(netloc=host[4:]).geturl() + return url + + +def is_google_drive_addon_url(url: str) -> bool: + """Return True if url is a Google Drive link (after normalizing www.).""" + return is_google_drive_url(_normalize_google_drive_url(url)) + + +def download_google_drive_zip(url: str, dest: Path) -> None: + """Download a publicly shared Google Drive zip to dest via gdown.""" + gdown.download(url=_normalize_google_drive_url(url), output=str(dest), quiet=True) + + +def _addon_zip_path_for_url(addon_url: str, addon_zip_dir: Path) -> Path: + normalized = _normalize_google_drive_url(addon_url) + if is_google_drive_url(normalized): + file_id, _ = parse_url(normalized) + if file_id: + return addon_zip_dir / f'gdrive_{file_id}.zip' + download_name = urlparse(addon_url).path.replace(os.path.sep, '_') + return addon_zip_dir / f'{download_name}.zip' + + +@app.task(bind=True, acks_late=True, ignore_result=True) +def upgrade_pipelines( + self: Task, + urls: List[str] = UPGRADE_JOB_DEFAULT_URLS, + force: bool = False, +): + """Install addons from zip files over HTTP (including Google Drive share links)""" + conf = Config() + context: dict = {} + manager: JobManager = patch_manager(self.job_manager) + if utils.check_canceled(self, context): + manager.updateStatus(JobStatus.CANCELED) + return + + gc: GirderClient = self.girder_client + # zipfiles to extract after download is complete + addons_to_update_update: List[Path] = [] + + for addon in urls: + zipfile_path = _addon_zip_path_for_url(addon, conf.addon_zip_path) + had_existing_zip = zipfile_path.exists() + try: + if not had_existing_zip or force: + manager.write(f'Downloading {addon} to {zipfile_path}\n') + if is_google_drive_addon_url(addon): + download_google_drive_zip(addon, zipfile_path) + else: + request.urlretrieve(addon, filename=zipfile_path) + else: + manager.write(f'Skipping download of {zipfile_path}\n') + addons_to_update_update.append(zipfile_path) + except Exception as exc: + logger.exception('Failed to download addon %s', addon) + manager.write(f'Failed to download {addon}: {exc}\nSkipping.\n') + if zipfile_path.exists() and not had_existing_zip: + zipfile_path.unlink(missing_ok=True) + if utils.check_canceled(self, context, force=False): + manager.updateStatus(JobStatus.CANCELED) + return + + # remove and recreate the existing addon pipeline directory + shutil.rmtree(conf.addon_extracted_path) + # Seed base pipelines from the VIAME image when available (GPU workers only). + if conf.viame_pipeline_path.exists(): + shutil.copytree(conf.viame_pipeline_path, conf.get_extracted_pipeline_path(missing_ok=True)) + # Extract zipfiles over newly copied files. Right now the zip archives + # MUST contain the pipeline subdir (e.g. configs/pipelines) in their + # internal structure. + for zipfile_path in addons_to_update_update: + manager.write(f'Extracting {zipfile_path} to {str(conf.addon_extracted_path)}\n') + z = zipfile.ZipFile(zipfile_path) + z.extractall(conf.addon_extracted_path) + + if utils.check_canceled(self, context): + # Remove everything + shutil.rmtree(conf.addon_extracted_path) + manager.updateStatus(JobStatus.CANCELED) + gc.put('dive_configuration/static_pipeline_configs', json=EMPTY_JOB_SCHEMA) + return + + # finally, crawl the new files and report results + summary = discover_configs(conf.get_extracted_pipeline_path()) + manager.write(str(summary)) + gc.put('dive_configuration/static_pipeline_configs', json=summary) + # get a list of files in the zip directory for the installed configuration listing + downloaded = [] + + # Iterate directory + for path in os.listdir(conf.addon_zip_path): + # check if current path is a file + if os.path.isfile(os.path.join(conf.addon_zip_path, path)): + downloaded.append(path) + print('Downloaded Files') + print(downloaded) + gc.put('dive_configuration/installed_addons', json={'downloaded': downloaded}) diff --git a/server/dive_tasks/viame_config.py b/server/dive_tasks/viame_config.py new file mode 100644 index 000000000..db531bc70 --- /dev/null +++ b/server/dive_tasks/viame_config.py @@ -0,0 +1,100 @@ +import os +from pathlib import Path +from typing import Dict + +from GPUtil import getGPUs + +from dive_tasks import utils +from dive_utils.types import AvailableJobSchema + +EMPTY_JOB_SCHEMA: AvailableJobSchema = { + 'pipelines': {}, + 'training': { + 'configs': [], + 'default': None, + }, + 'models': {}, +} + + +def get_gpu_environment() -> Dict[str, str]: + """Get environment variables for using CUDA enabled GPUs.""" + env = os.environ.copy() + + gpu_uuid = env.get("WORKER_GPU_UUID") + gpus = [gpu.id for gpu in getGPUs() if gpu.uuid == gpu_uuid] + + # Only set this env var if WORKER_GPU_UUID was supplied, + # and it matches an installed GPU + if gpus: + env["CUDA_VISIBLE_DEVICES"] = str(gpus[0]) + # Support for NOAA python3.10 means removing the local venv from the path + env["PATH"] = env.get("PATH").replace("/opt/dive/local/venv/bin", "") + return env + + +_VIAME_WORKER_QUEUES = frozenset({'pipelines', 'training'}) + + +def _worker_requires_viame_install() -> bool: + """ + Only pipeline/training workers need a local VIAME install. + + Default (``celery``), ``local``, and dev ``localworker`` processes do not; + they never call :class:`Config` today, but this keeps :meth:`Config.__init__` + safe if a task is misrouted. + """ + queues = os.environ.get('WORKER_WATCHING_QUEUES', '') + watched = {q.strip() for q in queues.split(',') if q.strip()} + return bool(watched & _VIAME_WORKER_QUEUES) + + +class Config: + def __init__(self): + self.gpu_process_env = get_gpu_environment() + self.viame_install_directory = os.environ.get( + 'VIAME_INSTALL_PATH', + '/opt/noaa/viame', + ) + self.addon_root_directory = os.environ.get( + 'ADDON_ROOT_DIR', + '/tmp/addons', + ) + self.kwiver_log_level = os.environ.get( + 'KWIVER_DEFAULT_LOG_LEVEL', + 'warn', + ) + + self.pipeline_subdir = 'configs/pipelines' + self.viame_install_path = Path(self.viame_install_directory) + self.viame_setup_script = self.viame_install_path / "setup_viame.sh" + self.viame_executable = self.viame_install_path / "bin" / "viame" + self.viame_pipeline_path = self.viame_install_path / self.pipeline_subdir + + if _worker_requires_viame_install(): + self.require_viame_install() + + self.addon_root_path = Path(self.addon_root_directory) + self.addon_zip_path = utils.make_directory(self.addon_root_path / 'zips') + self.addon_extracted_path = utils.make_directory(self.addon_root_path / 'extracted') + + # Set include directory to include pipelines from this path + # https://github.com/VIAME/VIAME/issues/131 + self.gpu_process_env['SPROKIT_PIPE_INCLUDE_PATH'] = str( + self.addon_extracted_path / self.pipeline_subdir + ) + + def require_viame_install(self) -> None: + assert self.viame_install_path.exists(), "VIAME Base install directory missing." + assert self.viame_setup_script.is_file(), "VIAME Setup Script missing" + assert self.viame_executable.is_file(), "VIAME Executable missing" + assert self.viame_pipeline_path.exists(), "VIAME common pipe directory missing." + + def get_extracted_pipeline_path(self, missing_ok=False) -> Path: + """ + Includes subdirectory for pipelines + """ + pipeline_path = self.addon_extracted_path / self.pipeline_subdir + if not missing_ok: + assert pipeline_path.exists(), f"Missing path {pipeline_path}" + return pipeline_path diff --git a/server/tests/test_google_drive_download.py b/server/tests/test_google_drive_download.py index cac7f2b12..f91bee3cb 100644 --- a/server/tests/test_google_drive_download.py +++ b/server/tests/test_google_drive_download.py @@ -3,7 +3,7 @@ import pytest -from dive_tasks.tasks import ( +from dive_tasks.upgrade_pipelines import ( _addon_zip_path_for_url, download_google_drive_zip, is_google_drive_addon_url, @@ -44,7 +44,7 @@ def test_addon_zip_path_keeps_http_path_naming(tmp_path: Path): def test_download_google_drive_zip_uses_gdown(tmp_path: Path): dest = tmp_path / 'addon.zip' url = 'https://www.drive.google.com/file/d/abc123XYZ/view?usp=sharing' - with patch('dive_tasks.tasks.gdown.download') as mock_download: + with patch('dive_tasks.upgrade_pipelines.gdown.download') as mock_download: download_google_drive_zip(url, dest) mock_download.assert_called_once_with( url='https://drive.google.com/file/d/abc123XYZ/view?usp=sharing', diff --git a/server/tests/test_inject_metadata_file.py b/server/tests/test_inject_metadata_file.py index 0b3e0d557..30f516d11 100644 --- a/server/tests/test_inject_metadata_file.py +++ b/server/tests/test_inject_metadata_file.py @@ -7,7 +7,7 @@ from pathlib import Path -from dive_tasks.tasks import _inject_dataset_metadata_file +from dive_tasks.run_pipeline import _inject_dataset_metadata_file class FakeGirderClient: diff --git a/server/tests/test_remote_ffprobe.py b/server/tests/test_remote_ffprobe.py index be79e1435..2b25d323c 100644 --- a/server/tests/test_remote_ffprobe.py +++ b/server/tests/test_remote_ffprobe.py @@ -232,7 +232,7 @@ def test_is_frame_misaligned_accepts_url_and_headers(): def _run_convert_video(task, **kwargs): """Invoke convert_video's underlying function with an explicit bound self.""" - from dive_tasks.tasks import convert_video + from dive_tasks.convert_video import convert_video # PromiseProxy.__wrapped__ is a bound method; call the unbound function. return convert_video.__wrapped__.__func__(task, **kwargs) @@ -264,10 +264,10 @@ def test_convert_video_skip_path_avoids_download(): task.job_manager = MagicMock() with ( - patch('dive_tasks.tasks.patch_manager') as patch_mgr, - patch('dive_tasks.tasks.utils.ffprobe_format_and_streams', return_value=probe), - patch('dive_tasks.tasks.is_frame_misaligned', return_value=False), - patch('dive_tasks.tasks.resolve_annotation_fps', return_value=30.0), + patch('dive_tasks.convert_video.patch_manager') as patch_mgr, + patch('dive_tasks.convert_video.utils.ffprobe_format_and_streams', return_value=probe), + patch('dive_tasks.convert_video.is_frame_misaligned', return_value=False), + patch('dive_tasks.convert_video.resolve_annotation_fps', return_value=30.0), ): patch_mgr.return_value = MagicMock() _run_convert_video( @@ -315,15 +315,15 @@ def test_convert_video_downloads_when_transcode_required(): task.job_manager = MagicMock() with ( - patch('dive_tasks.tasks.patch_manager') as patch_mgr, - patch('dive_tasks.tasks.utils.ffprobe_format_and_streams', return_value=probe), - patch('dive_tasks.tasks.is_frame_misaligned', return_value=False), - patch('dive_tasks.tasks.utils.stream_subprocess'), + patch('dive_tasks.convert_video.patch_manager') as patch_mgr, + patch('dive_tasks.convert_video.utils.ffprobe_format_and_streams', return_value=probe), + patch('dive_tasks.convert_video.is_frame_misaligned', return_value=False), + patch('dive_tasks.convert_video.utils.stream_subprocess'), patch( - 'dive_tasks.tasks.check_and_fix_frame_alignment', + 'dive_tasks.convert_video.check_and_fix_frame_alignment', side_effect=lambda task, path, context, manager: path, ), - patch('dive_tasks.tasks.resolve_annotation_fps', return_value=25.0), + patch('dive_tasks.convert_video.resolve_annotation_fps', return_value=25.0), ): patch_mgr.return_value = MagicMock() _run_convert_video( @@ -352,9 +352,9 @@ def test_convert_video_cancel_during_remote_probe_does_not_download(): task.job_manager = MagicMock() with ( - patch('dive_tasks.tasks.patch_manager') as patch_mgr, + patch('dive_tasks.convert_video.patch_manager') as patch_mgr, patch( - 'dive_tasks.tasks.utils.ffprobe_format_and_streams', + 'dive_tasks.convert_video.utils.ffprobe_format_and_streams', side_effect=CanceledError('Job was canceled'), ), ): @@ -398,10 +398,10 @@ def test_convert_video_cancel_during_remote_alignment_does_not_download(): task.job_manager = MagicMock() with ( - patch('dive_tasks.tasks.patch_manager') as patch_mgr, - patch('dive_tasks.tasks.utils.ffprobe_format_and_streams', return_value=probe), + patch('dive_tasks.convert_video.patch_manager') as patch_mgr, + patch('dive_tasks.convert_video.utils.ffprobe_format_and_streams', return_value=probe), patch( - 'dive_tasks.tasks.is_frame_misaligned', + 'dive_tasks.convert_video.is_frame_misaligned', side_effect=CanceledError('Job was canceled'), ), ):