From ab317c3ee791fe61e16c971714ced55523805c51 Mon Sep 17 00:00:00 2001 From: ljcornel Date: Tue, 28 May 2024 11:49:28 +0200 Subject: [PATCH] Undo changes to video tempfile handling, register atexit handler for tempfile deletion --- geti_sdk/data_models/media.py | 9 ++++--- .../rest_clients/media_client/video_client.py | 25 ++++++++++++------- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/geti_sdk/data_models/media.py b/geti_sdk/data_models/media.py index 84e4c054..47530efe 100644 --- a/geti_sdk/data_models/media.py +++ b/geti_sdk/data_models/media.py @@ -15,6 +15,7 @@ import abc import os import tempfile +from multiprocessing import Lock from pprint import pformat from typing import Any, Dict, List, Optional, Union @@ -308,6 +309,7 @@ def __attrs_post_init__(self): """ self._data: Optional[str] = None self._needs_tempfile_deletion: bool = False + self._data_lock: Lock = Lock() @property def identifier(self) -> VideoIdentifier: @@ -396,9 +398,10 @@ def __del__(self): Clean up the temporary file created to store the video data (if any). This method is called when the Video object is deleted. """ - if self._needs_tempfile_deletion: - if os.path.exists(self._data): - os.remove(self._data) + with self._data_lock: + if self._needs_tempfile_deletion: + if os.path.exists(self._data) and os.path.isfile(self._data): + os.remove(self._data) @attr.define(slots=False) diff --git a/geti_sdk/rest_clients/media_client/video_client.py b/geti_sdk/rest_clients/media_client/video_client.py index dd0acb80..e61b1b08 100644 --- a/geti_sdk/rest_clients/media_client/video_client.py +++ b/geti_sdk/rest_clients/media_client/video_client.py @@ -11,9 +11,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions # and limitations under the License. - - +import atexit import os +import tempfile from datetime import datetime from typing import Optional, Sequence, Union @@ -80,7 +80,12 @@ def upload_video( f"shape {video.shape}" ) from error timestamp = datetime.now().strftime("%Y%m%d%H%M%S") - video_path = f"temp_video_{timestamp}.avi" + video_file = tempfile.NamedTemporaryFile( + prefix="geti-sdk_temp_video_", suffix=f"_{timestamp}.avi", delete=False + ) + # Close the file, opencv will open it again from the path + video_file.close() + video_path = video_file.name out = cv2.VideoWriter( video_path, cv2.VideoWriter_fourcc("M", "J", "P", "G"), @@ -98,13 +103,15 @@ def upload_video( uploaded_video = MediaRESTConverter.from_dict( input_dict=video_dict, media_type=Video ) - # Clean up temp file + uploaded_video._data = video_path if temporary_file_created: - # We do not keep the video file in case of uploading a numpy array representation. - # The intuition is that the user is responsible for managing the original video. - os.remove(video_path) - else: - uploaded_video._data = video_path + uploaded_video._needs_tempfile_deletion = True + + # Register cleanup function on system exit to ensure __del__ gets called + def clean_temp_video(): + uploaded_video.__del__() + + atexit.register(clean_temp_video) return uploaded_video def upload_folder(