Skip to content

Commit

Permalink
Undo changes to video tempfile handling, register atexit handler for …
Browse files Browse the repository at this point in the history
…tempfile deletion
  • Loading branch information
ljcornel committed May 28, 2024
1 parent 8e77344 commit ab317c3
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 12 deletions.
9 changes: 6 additions & 3 deletions geti_sdk/data_models/media.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
25 changes: 16 additions & 9 deletions geti_sdk/rest_clients/media_client/video_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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"),
Expand All @@ -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(
Expand Down

0 comments on commit ab317c3

Please sign in to comment.