diff --git a/.gitignore b/.gitignore index 2ccc2825..97b66483 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,5 @@ Examples/reprocessDataServer.py newsletter.py .DS_Store -.vscode/ \ No newline at end of file +.vscode/ +sandbox_process_trial.py diff --git a/main.py b/main.py index d5994224..bff335b0 100644 --- a/main.py +++ b/main.py @@ -26,6 +26,7 @@ from utilsChecker import writeTRCfrom3DKeypoints from utilsChecker import popNeutralPoseImages from utilsChecker import rotateIntrinsics +from utilsChecker import loadLidarIntrinsicMatrix from utilsSync import synchronizeVideos from utilsDetector import runPoseDetector from utilsAugmenter import augmentTRC @@ -217,10 +218,13 @@ def main(sessionName, trialName, trial_id, cameras_to_use=['all'], loadedCamParams = {} for camName in cameraDirectories: camDir = cameraDirectories[camName] + lidarIntrinsicPath = os.path.join(camDir, 'InputMedia', trialName, + 'camera_matrix.csv') + hasLidarIntrinsics = os.path.exists(lidarIntrinsicPath) # Intrinsics ###################################################### # Intrinsics and extrinsics already exist for this session. if os.path.exists( - os.path.join(camDir,"cameraIntrinsicsExtrinsics.pickle")): + os.path.join(camDir,"cameraIntrinsicsExtrinsics.pickle")) and not hasLidarIntrinsics: logging.info("Load extrinsics for {} - already existing".format( camName)) CamParams = loadCameraParameters( @@ -240,7 +244,10 @@ def main(sessionName, trialName, trial_id, cameras_to_use=['all'], if os.path.exists(permIntrinsicDir): CamParams = loadCameraParameters( os.path.join(permIntrinsicDir, - 'cameraIntrinsics.pickle')) + 'cameraIntrinsics.pickle')) + if hasLidarIntrinsics: + CamParams['intrinsicMat'] = loadLidarIntrinsicMatrix(lidarIntrinsicPath) + logging.info("Using LiDAR camera_matrix.csv intrinsics for {}".format(camName)) # Intrinsics do not exist throw an error. Eventually the # webapp will give you the opportunity to compute them. diff --git a/requirements.txt b/requirements.txt index 0937e2f6..7837d583 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,3 +17,4 @@ ffmpeg-python psutil boto3 pytest +zstandard diff --git a/sandbox_app_local.py b/sandbox_app_local.py new file mode 100644 index 00000000..f93bf64c --- /dev/null +++ b/sandbox_app_local.py @@ -0,0 +1,108 @@ +import argparse +import logging +import time +import traceback + +import requests + +from utils import getTrialJson, makeRequestWithRetry +from utilsAPI import getAPIURL +from utilsAuth import getToken +from utilsServer import processTrial + + +logging.basicConfig(format="[%(asctime)s] [%(levelname)s] %(message)s", + level=logging.INFO, + datefmt="%Y-%m-%d %H:%M:%S") + +API_URL = getAPIURL() +API_TOKEN = getToken() + + +def get_trial_type(trial): + if trial["name"] == "calibration": + return "calibration" + if trial["name"] == "neutral": + return "static" + return "dynamic" + + +def get_next_trial(worker_type="all"): + queue_path = "trials/dequeue/?workerType={}".format(worker_type) + response = requests.get( + "{}{}".format(API_URL, queue_path), + headers={"Authorization": "Token {}".format(API_TOKEN)}) + + if response.status_code == 404: + return None + + response.raise_for_status() + return response.json() + + +def patch_trial_status(trial_id, status): + trial_url = "{}{}{}/".format(API_URL, "trials/", trial_id) + makeRequestWithRetry( + "PATCH", + trial_url, + data={"status": status}, + headers={"Authorization": "Token {}".format(API_TOKEN)}) + + +def process_local_trial(trial, skip_status_patch=False): + trial_type = get_trial_type(trial) + logging.info("Processing trial %s from session %s as %s.", + trial["id"], trial["session"], trial_type) + + try: + processTrial(trial["session"], trial["id"], trial_type=trial_type, + isDocker=False, deleteLocalFolder=False) + except Exception: + traceback.print_exc() + if not skip_status_patch: + patch_trial_status(trial["id"], "error") + raise + + if not skip_status_patch: + patch_trial_status(trial["id"], "done") + + logging.info("Finished trial %s.", trial["id"]) + + +def main(): + parser = argparse.ArgumentParser( + description="Local polling trial processor for sandbox testing.") + parser.add_argument("--trial-id", help="Process a specific trial once instead of polling.") + parser.add_argument("--worker-type", default="all", + help="Worker type to use when dequeuing. Defaults to all.") + parser.add_argument("--sleep-seconds", type=float, default=2, + help="Seconds to sleep when no queued trial is found.") + parser.add_argument("--skip-status-patch", action="store_true", + help="Do not PATCH the trial status to done/error.") + args = parser.parse_args() + + if args.trial_id: + process_local_trial(getTrialJson(args.trial_id), args.skip_status_patch) + return + + while True: + try: + trial = get_next_trial(args.worker_type) + except Exception: + traceback.print_exc() + time.sleep(args.sleep_seconds) + continue + + if trial is None: + logging.info("No queued trial found for workerType=%s.", args.worker_type) + time.sleep(args.sleep_seconds) + continue + + try: + process_local_trial(trial, args.skip_status_patch) + except Exception: + logging.info("Trial %s failed. Continuing local polling loop.", trial["id"]) + + +if __name__ == "__main__": + main() diff --git a/utils.py b/utils.py index c7ee8d03..962aa74b 100644 --- a/utils.py +++ b/utils.py @@ -10,9 +10,13 @@ import glob import mimetypes import subprocess +import struct import zipfile import time import datetime +import tempfile +import lzma +import zlib import numpy as np import pandas as pd @@ -24,6 +28,9 @@ API_URL = getAPIURL() API_TOKEN = getToken() +DEPTH_DB_LEVEL = 10 +DEPTH_DB_TRANSFORM = "vertical_delta_shuffle16" +DEPTH_CONTAINER_MAGIC = b"OCDEPTHDB1\n" #%% Rest of utils @@ -102,7 +109,316 @@ def importMetadata(filePath): def download_file(url, file_name): with urllib.request.urlopen(url) as response, open(file_name, 'wb') as out_file: shutil.copyfileobj(response, out_file) - + +def uploadFileToS3(filePath): + data = {'fileName':os.path.split(filePath)[1]} + response = makeRequestWithRetry('GET', + API_URL + "sessions/null/get_presigned_url/", + data=data) + r = response.json() + + with open(filePath, 'rb') as file: + files = {'file': file} + makeRequestWithRetry('POST', + r['url'], + data=r['fields'], + files=files) + + return r['fields']['key'] + +def hasLidarData(video): + is_lidar = video.get('isLidar', False) + + if isinstance(is_lidar, str): + return is_lidar.lower() == 'true' + + return bool(is_lidar) + +def _safeExtractZip(zip_path, extract_dir): + with zipfile.ZipFile(zip_path, 'r') as zip_ref: + extract_root = os.path.abspath(extract_dir) + for member in zip_ref.infolist(): + member_path = os.path.abspath(os.path.join(extract_dir, member.filename)) + if os.path.commonpath([extract_root, member_path]) != extract_root: + raise Exception('Unsafe path in zip file: {}'.format(member.filename)) + zip_ref.extractall(extract_dir) + +def _findSingleMovFile(folder): + mov_files = glob.glob(os.path.join(folder, '**', '*.mov'), recursive=True) + mov_files += glob.glob(os.path.join(folder, '**', '*.MOV'), recursive=True) + mov_files = list(dict.fromkeys(mov_files)) + + if len(mov_files) == 0: + raise Exception('No .mov file found in LiDAR zip bundle.') + if len(mov_files) > 1: + raise Exception('Multiple .mov files found in LiDAR zip bundle: {}'.format(mov_files)) + + return mov_files[0] + +def _zipFolder(folder, zip_path): + with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: + for root, dirs, files in os.walk(folder): + for file in files: + file_path = os.path.join(root, file) + zipf.write(file_path, os.path.relpath(file_path, folder)) + +def _decompressDepthBlob(blob, compression): + if compression in (None, "", "none"): + return blob + if compression in ("deflate_raw", "zlib"): + return zlib.decompress(blob, wbits=-15) + if compression == "zlib_wrapped": + return zlib.decompress(blob) + if compression == "lzma": + return lzma.decompress(blob) + raise ValueError("Unsupported depth compression: {}".format(compression)) + +def _loadDepthFrames(metadata_path): + with open(metadata_path, "r", encoding="utf-8") as f: + meta = json.load(f) + + if meta.get("format") != "uint16": + raise ValueError("Unsupported depth format: {}".format(meta.get("format"))) + if meta.get("byte_order") != "little_endian": + raise ValueError("Unsupported depth byte_order: {}".format(meta.get("byte_order"))) + + width = int(meta["width"]) + height = int(meta["height"]) + frame_count = int(meta["frame_count"]) + bytes_per_pixel = int(meta.get("bytes_per_pixel", 2)) + frame_byte_count = int(meta.get("frame_byte_count", width * height * bytes_per_pixel)) + if bytes_per_pixel != 2 or frame_byte_count != width * height * 2: + raise ValueError("Unsupported depth frame byte count.") + + depth_path = os.path.join(os.path.dirname(metadata_path), meta.get("file", "depth.bin")) + blob = open(depth_path, "rb").read() + + if meta.get("frame_layout") == "length_prefixed": + frames = [] + pos = 0 + for frame_index in range(frame_count): + if pos + 4 > len(blob): + raise ValueError("Unexpected EOF reading depth frame {} length".format(frame_index)) + compressed_len = struct.unpack_from(" len(blob): + raise ValueError("Unexpected EOF reading depth frame {} payload".format(frame_index)) + frame_bytes = _decompressDepthBlob(blob[pos:end], meta.get("compression", "none")) + pos = end + if len(frame_bytes) != frame_byte_count: + raise ValueError("Depth frame {} decoded to {} bytes; expected {}".format( + frame_index, len(frame_bytes), frame_byte_count)) + frames.append(np.frombuffer(frame_bytes, dtype=" 1: + raise Exception("Multiple camera_matrix.csv files found in LiDAR zip bundle: {}".format(csv_paths)) + return csv_paths[0] + +def _copyLidarCameraMatrixCsv(extract_dir, video_path): + csv_path = _findSingleCameraMatrixCsv(extract_dir) + if csv_path is None: + print("No camera_matrix.csv found in LiDAR zip bundle.") + return None + + output_path = os.path.join(os.path.dirname(video_path), "camera_matrix.csv") + shutil.copyfile(csv_path, output_path) + return output_path + +def _downloadLidarResultCameraMatrixCsv(video, video_path, trial_id): + trial = getTrialJson(trial_id) + matching_results = [ + result for result in trial["results"] + if result["tag"] == "lidar_data" + and result.get("device_id") == video.get("device_id") + and result.get("media") + ] + if len(matching_results) == 0: + print("No lidar_data trial result found for device_id {}.".format(video.get("device_id"))) + return None + + zip_path = os.path.splitext(video_path)[0] + "_lidar_data.zip" + extract_dir = os.path.splitext(zip_path)[0] + + try: + download_file(matching_results[0]["media"], zip_path) + if not zipfile.is_zipfile(zip_path): + print("lidar_data result for device_id {} is not a zip file.".format(video.get("device_id"))) + return None + if os.path.exists(extract_dir): + shutil.rmtree(extract_dir) + os.makedirs(extract_dir, exist_ok=True) + _safeExtractZip(zip_path, extract_dir) + return _copyLidarCameraMatrixCsv(extract_dir, video_path) + finally: + if os.path.exists(extract_dir): + shutil.rmtree(extract_dir) + if os.path.exists(zip_path): + os.remove(zip_path) + +def updateVideoFileInAPI(video, video_path): + if 'id' not in video: + raise Exception('Cannot update video file in API because the video JSON has no id.') + + media_url = uploadFileToS3(video_path) + video_url = "{}{}{}/".format(API_URL, "videos/", video['id']) + headers = {"Authorization": "Token {}".format(API_TOKEN)} + patch_options = [{'video_url': media_url}, + {'media_url': media_url}, + {'video': media_url}] + + last_error = None + for data in patch_options: + try: + rVideo = makeRequestWithRetry('PATCH', + video_url, + data=data, + headers=headers) + try: + video_response = rVideo.json() + video['video'] = video_response.get('video', media_url) + except ValueError: + video['video'] = media_url + return media_url + except requests.exceptions.HTTPError as e: + if e.response is not None and e.response.status_code == 400: + last_error = e + continue + raise + + raise last_error + +def downloadVideoFile(video, video_path, trial_id, use_lidar_intrinsics=False): + if not hasLidarData(video): + download_file(video["video"], video_path) + return + + zip_path = os.path.splitext(video_path)[0] + '.zip' + extract_dir = os.path.splitext(zip_path)[0] + + download_file(video["video"], zip_path) + + if not zipfile.is_zipfile(zip_path): + if os.path.exists(video_path): + os.remove(video_path) + shutil.move(zip_path, video_path) + if use_lidar_intrinsics: + _downloadLidarResultCameraMatrixCsv(video, video_path, trial_id) + return + + if os.path.exists(extract_dir): + shutil.rmtree(extract_dir) + os.makedirs(extract_dir, exist_ok=True) + + _safeExtractZip(zip_path, extract_dir) + mov_path = _findSingleMovFile(extract_dir) + + if os.path.exists(video_path): + os.remove(video_path) + shutil.move(mov_path, video_path) + if use_lidar_intrinsics: + _copyLidarCameraMatrixCsv(extract_dir, video_path) + _compressLidarDepthFiles(extract_dir) + + updateVideoFileInAPI(video, video_path) + + if os.path.exists(zip_path): + os.remove(zip_path) + _zipFolder(extract_dir, zip_path) + + deleteResult(trial_id, tag='lidar_data', device_id=video.get('device_id')) + postFileToTrial(zip_path,trial_id,tag='lidar_data', + device_id=video.get('device_id')) + + shutil.rmtree(extract_dir) + os.remove(zip_path) + +def downloadUnmappedLidarVideoFile(video, trial_id): + temp_dir = tempfile.mkdtemp(prefix='opencap_lidar_') + video_path = os.path.join(temp_dir, '{}.mov'.format(trial_id)) + + try: + downloadVideoFile(video, video_path, trial_id) + finally: + if os.path.exists(temp_dir): + shutil.rmtree(temp_dir) + def getTrialJson(trial_id): response = makeRequestWithRetry('GET', API_URL + "trials/{}/".format(trial_id), @@ -229,21 +545,36 @@ def downloadVideosFromServer(session_id,trial_id, isDocker=True, for k, video in enumerate(trial["videos"]): os.makedirs(os.path.join(session_path, "Videos", "Cam{}".format(k), "InputMedia", trial_name), exist_ok=True) video_path = os.path.join(session_path, "Videos", "Cam{}".format(k), "InputMedia", trial_name, trial_id + ".mov") - download_file(video["video"], video_path) + downloadVideoFile(video, video_path, trial_id, + use_lidar_intrinsics=isCalibration) mappingCamDevice[video["device_id"].replace('-', '').upper()] = k with open(os.path.join(session_path, "Videos", 'mappingCamDevice.pickle'), 'wb') as handle: pickle.dump(mappingCamDevice, handle) else: with open(os.path.join(session_path, "Videos", 'mappingCamDevice.pickle'), 'rb') as handle: mappingCamDevice = pickle.load(handle) - for video in trial["videos"]: - k = mappingCamDevice[video["device_id"].replace('-', '').upper()] + for video in trial["videos"]: + device_id = video["device_id"].replace('-', '').upper() + if device_id not in mappingCamDevice: + if not isCalibration: + if hasLidarData(video): + print('Processing unmapped LiDAR video with device_id {} outside the local Videos folder.'.format(video["device_id"])) + downloadUnmappedLidarVideoFile(video, trial_id) + else: + print('Skipping video with device_id {} because it was not in the calibration camera mapping.'.format(video["device_id"])) + continue + raise KeyError(device_id) + k = mappingCamDevice[device_id] videoDir = os.path.join(session_path, "Videos", "Cam{}".format(k), "InputMedia", trial_name) os.makedirs(videoDir, exist_ok=True) video_path = os.path.join(videoDir, trial_id + ".mov") if not os.path.exists(video_path): if video['video'] : - download_file(video["video"], video_path) + downloadVideoFile(video, video_path, trial_id, + use_lidar_intrinsics=isCalibration) + elif (isCalibration and hasLidarData(video) and + not os.path.exists(os.path.join(videoDir, "camera_matrix.csv"))): + _downloadLidarResultCameraMatrixCsv(video, video_path, trial_id) # Import and save metadata sessionYamlPath = os.path.join(session_path, "sessionMetadata.yaml") @@ -453,13 +784,14 @@ def getMetadataFromServer(session_id,justCheckerParams=False): return session_desc -def deleteResult(trial_id, tag=None,resultNum=None): +def deleteResult(trial_id, tag=None,resultNum=None,device_id=None): # Delete specific result number, or all results with a specific tag, or all results if tag==None if resultNum != None: resultNums = [resultNum] elif tag != None: trial = getTrialJson(trial_id) - resultNums = [r['id'] for r in trial['results'] if r['tag']==tag] + resultNums = [r['id'] for r in trial['results'] if r['tag']==tag + and (device_id is None or r.get('device_id')==device_id)] elif tag == None: trial = getTrialJson(trial_id) @@ -908,28 +1240,14 @@ def getModelAndMetadata(session_id,session_path,simplePath=False): return def postFileToTrial(filePath,trial_id,tag,device_id): - - # get S3 link - data = {'fileName':os.path.split(filePath)[1]} - response = makeRequestWithRetry('GET', - API_URL + "sessions/null/get_presigned_url/", - data=data) - r = response.json() - - # upload to S3 - files = {'file': open(filePath, 'rb')} - makeRequestWithRetry('POST', - r['url'], - data=r['fields'], - files=files) - files["file"].close() + media_url = uploadFileToS3(filePath) # post link to and data to results data = { "trial": trial_id, "tag": tag, "device_id" : device_id, - "media_url" : r['fields']['key'] + "media_url" : media_url } rResult = makeRequestWithRetry('POST', @@ -938,7 +1256,7 @@ def postFileToTrial(filePath,trial_id,tag,device_id): headers = {"Authorization": "Token {}".format(API_TOKEN)}) if rResult.status_code != 201: - print('server response was + ' + str(r.status_code)) + print('server response was + ' + str(rResult.status_code)) else: print('Result posted to S3.') diff --git a/utilsChecker.py b/utilsChecker.py index 4d464e54..00014476 100644 --- a/utilsChecker.py +++ b/utilsChecker.py @@ -284,6 +284,15 @@ def saveCameraParameters(filename,CameraParams): return True +def loadLidarIntrinsicMatrix(cameraMatrixPath): + intrinsicMat = np.loadtxt(cameraMatrixPath, delimiter=',') + if intrinsicMat.shape != (3, 3): + raise Exception("LiDAR camera_matrix.csv must contain a 3x3 matrix.") + + intrinsicMat = intrinsicMat.astype(np.float64) + intrinsicMat[0, 2], intrinsicMat[1, 2] = intrinsicMat[1, 2], intrinsicMat[0, 2] + return intrinsicMat + #%% def getVideoRotation(videoPath): @@ -420,6 +429,19 @@ def calcExtrinsics(imageFileName, CameraParams, CheckerBoardParams, grayColor, CheckerBoardParams['dimensions'], cv2.CALIB_CB_ADAPTIVE_THRESH) + # Fallback: if the standard detector fails, try the SB variant (more robust to + # certain lighting/contrast conditions where adaptive thresholding struggles). + # Using ACCURACY|LARGER without EXHAUSTIVE to keep the fallback reasonably fast. + corners2_from_sb = False + if not ret: + ret_sb, corners_sb, _ = cv2.findChessboardCornersSBWithMeta( + grayColor, CheckerBoardParams['dimensions'], + cv2.CALIB_CB_ACCURACY | cv2.CALIB_CB_LARGER) + if ret_sb: + ret = True + corners = corners_sb + corners2_from_sb = True + # If desired number of corners can be detected then, # refine the pixel coordinates and display # them on the images of checker board @@ -429,8 +451,12 @@ def calcExtrinsics(imageFileName, CameraParams, CheckerBoardParams, # Refining pixel coordinates # for given 2d points. - corners2 = cv2.cornerSubPix( - grayColor, corners, (11, 11), (-1, -1), criteria) / imageUpsampleFactor + # SBWithMeta corners are already subpixel-accurate; skip cornerSubPix in that case. + if corners2_from_sb: + corners2 = corners / imageUpsampleFactor + else: + corners2 = cv2.cornerSubPix( + grayColor, corners, (11, 11), (-1, -1), criteria) / imageUpsampleFactor twodpoints.append(corners2)