Fix copy_images - video

Author: dshabin

superannotate/db/project_images.py

@@ -17,7 +17,8 @@
 from .project_api import get_project_and_folder_metadata
 from .projects import (
     __create_image, get_image_array_to_upload,
-    get_project_default_image_quality_in_editor, upload_image_array_to_s3
+    get_project_default_image_quality_in_editor, upload_image_array_to_s3,
+    _get_available_image_counts
 )
 from .utils import _get_upload_auth_token
 
@@ -99,12 +100,9 @@ def upload_image_to_project(
         folder_id = get_project_root_folder_id(project)
 
     team_id, project_id = project["team_id"], project["id"]
-    params = {
-        'team_id': team_id,
-        'folder_id' : folder_id
-    }
-    res = _get_upload_auth_token(params=params,project_id=project_id)
-    prefix = res['filePath']    
+    params = {'team_id': team_id, 'folder_id': folder_id}
+    res = _get_upload_auth_token(params=params, project_id=project_id)
+    prefix = res['filePath']
     s3_session = boto3.Session(
         aws_access_key_id=res['accessKeyId'],
         aws_secret_access_key=res['secretAccessKey'],
@@ -224,18 +222,23 @@ def copy_images(
     )
     if image_names is None:
         image_names = search_images((source_project, source_project_folder))
+
+    limit = _get_available_image_counts(
+        destination_project, destination_project_folder
+    )
+    image_names = image_names[:limit]
     res = _copy_images(
         (source_project, source_project_folder),
         (destination_project, destination_project_folder), image_names,
         include_annotations, copy_annotation_status, copy_pin
     )
     logger.info(
         "Copied images %s from %s to %s. Number of skipped images %s",
-        image_names,
-        source_project["name"] + "" if source_project_folder is None else "/" +
-        source_project_folder["name"], destination_project["name"] +
-        "" if destination_project_folder is None else destination_project["name"] + "/" +
-        destination_project_folder["name"], len(res["skipped"])
+        image_names, source_project["name"] +
+        "" if source_project_folder is None else source_project["name"] + "/" +
+        source_project_folder["name"], destination_project["name"] + ""
+        if destination_project_folder is None else destination_project["name"] +
+        "/" + destination_project_folder["name"], len(res["skipped"])
     )
     return res["skipped"]
 

superannotate/db/projects.py

@@ -221,7 +221,7 @@ def _get_video_frames_count(video_path):
     return total_num_of_frames
 
 
-def _get_video_fps_ration(target_fps,video,ratio):
+def _get_video_fps_ration(target_fps, video, ratio):
     """
     Get video fps / target fps ratio
     """
@@ -230,7 +230,7 @@ def _get_video_fps_ration(target_fps,video,ratio):
         logger.warning(
             "Video frame rate %s smaller than target frame rate %s. Cannot change frame rate.",
             video_fps, target_fps
-            )
+        )
     else:
         logger.info(
             "Changing video frame rate from %s to target frame rate %s.",
@@ -239,22 +239,24 @@ def _get_video_fps_ration(target_fps,video,ratio):
         ratio = video_fps / target_fps
     return ratio
 
-def _get_available_image_counts(project,folder):
+
+def _get_available_image_counts(project, folder):
     if folder:
         folder_id = folder["id"]
     else:
         folder_id = get_project_root_folder_id(project)
-    params = {'team_id': project['team_id'] , 'folder_id' : folder_id }
-    res = _get_upload_auth_token(params=params,project_id=project['id'])
+    params = {'team_id': project['team_id'], 'folder_id': folder_id}
+    res = _get_upload_auth_token(params=params, project_id=project['id'])
     return res['availableImageCount']
 
+
 def _get_video_rotate_code(video_path):
     rotate_code = None
     try:
         cv2_rotations = {
-            90 : cv2.ROTATE_90_CLOCKWISE,
-            180 : cv2.ROTATE_180,
-            270 :cv2.ROTATE_90_COUNTERCLOCKWISE,
+            90: cv2.ROTATE_90_CLOCKWISE,
+            180: cv2.ROTATE_180,
+            270: cv2.ROTATE_90_COUNTERCLOCKWISE,
         }
 
         meta_dict = ffmpeg.probe(str(video_path))
@@ -276,14 +278,26 @@ def _get_video_rotate_code(video_path):
     return rotate_code
 
 
-def _extract_frames_from_video(start_time,end_time,ratio,video,video_path,tempdir,limit,rotate_code,total_num_of_frames):
+def _extract_frames_from_video(
+    start_time, end_time, video_path, tempdir, limit, target_fps
+):
+    video = cv2.VideoCapture(str(video_path), cv2.CAP_FFMPEG)
+    if not video.isOpened():
+        raise SABaseException(0, "Couldn't open video file " + str(video_path))
+    total_num_of_frames = _get_video_frames_count(video_path)
+    logger.info("Video frame count is %s.", total_num_of_frames)
+    ratio = 1.0
+    if target_fps:
+        ratio = _get_video_fps_ration(target_fps, video, ratio)
+    rotate_code = _get_video_rotate_code(video_path)
     video_name = Path(video_path).stem
     frame_no = 0
     frame_no_with_change = 1.0
     extracted_frame_no = 1
     logger.info("Extracting frames from video to %s.", tempdir.name)
     zero_fill_count = len(str(total_num_of_frames))
-    while extracted_frame_no < (limit + 1) :
+    extracted_frames_paths = []
+    while len(extracted_frames_paths) < limit:
         success, frame = video.read()
         if not success:
             break
@@ -298,16 +312,16 @@ def _extract_frames_from_video(start_time,end_time,ratio,video,video_path,tempdi
             continue
         if rotate_code:
             frame = cv2.rotate(frame, rotate_code)
-        cv2.imwrite(
-            str(
-                Path(tempdir.name) / (
-                    video_name + "_" +
-                    str(extracted_frame_no).zfill(zero_fill_count) + ".jpg"
-                )
-            ), frame
+        path = str(
+            Path(tempdir.name) / (
+                video_name + "_" +
+                str(extracted_frame_no).zfill(zero_fill_count) + ".jpg"
+            )
         )
+        cv2.imwrite(path, frame)
+        extracted_frames_paths.append(path)
         extracted_frame_no += 1
-    return extracted_frame_no - 1
+    return extracted_frames_paths
 
 
 def upload_video_to_project(
@@ -345,7 +359,7 @@ def upload_video_to_project(
     """
 
     project, folder = get_project_and_folder_metadata(project)
-    limit = _get_available_image_counts(project,folder)
+    limit = _get_available_image_counts(project, folder)
 
     upload_state = common.upload_state_int_to_str(project.get("upload_state"))
     if upload_state == "External":
@@ -354,23 +368,13 @@ def upload_video_to_project(
             "The function does not support projects containing images attached with URLs"
         )
     logger.info("Uploading from video %s.", str(video_path))
-    rotate_code = _get_video_rotate_code(video_path)
-    video = cv2.VideoCapture(str(video_path), cv2.CAP_FFMPEG)
-    if not video.isOpened():
-        raise SABaseException(0, "Couldn't open video file " + str(video_path))
-
-    total_num_of_frames = _get_video_frames_count(video_path)
-    logger.info("Video frame count is %s.", total_num_of_frames)
-    ratio = 1.0
-    if target_fps:
-        ratio = _get_video_fps_ration(target_fps,video,ratio)
     tempdir = tempfile.TemporaryDirectory()
-    extracted_frame_no = _extract_frames_from_video(start_time,end_time,ratio,
-                                                    video,video_path,tempdir,
-                                                    limit,rotate_code,total_num_of_frames)
+    extracted_frames = _extract_frames_from_video(
+        start_time, end_time, video_path, tempdir, limit, target_fps
+    )
     logger.info(
         "Extracted %s frames from video. Now uploading to platform.",
-        extracted_frame_no
+        len(extracted_frames)
     )
     filenames = upload_images_from_folder_to_project(
         (project, folder),
@@ -853,9 +857,7 @@ def upload_images_to_project(
     :rtype: tuple (3 members) of list of strs
     """
     project, folder = get_project_and_folder_metadata(project)
-    folder_name = project["name"] + (
-        f'/{folder["name"]}' if folder else ""
-    )
+    folder_name = project["name"] + (f'/{folder["name"]}' if folder else "")
     upload_state = common.upload_state_int_to_str(project.get("upload_state"))
     if upload_state == "External":
         raise SABaseException(
@@ -895,24 +897,23 @@ def upload_images_to_project(
     if len_img_paths == 0:
         return ([], [], duplicate_images)
 
-    
     if folder:
         folder_id = folder["id"]
     else:
         folder_id = get_project_root_folder_id(project)
 
-    params = {'team_id': team_id , 'folder_id' : folder_id }
+    params = {'team_id': team_id, 'folder_id': folder_id}
     uploaded = [[] for _ in range(_NUM_THREADS)]
     tried_upload = [[] for _ in range(_NUM_THREADS)]
     couldnt_upload = [[] for _ in range(_NUM_THREADS)]
     finish_event = threading.Event()
 
-    res = _get_upload_auth_token(params=params,project_id=project_id)
+    res = _get_upload_auth_token(params=params, project_id=project_id)
 
     prefix = res['filePath']
     limit = res['availableImageCount']
     images_to_upload = img_paths[:limit]
-    images_to_skip = img_paths[limit:]
+    images_to_skip = [str(path) for path in img_paths[limit:]]
     chunksize = int(math.ceil(len(images_to_upload) / _NUM_THREADS))
 
     tqdm_thread = threading.Thread(
@@ -927,8 +928,8 @@ def upload_images_to_project(
         t = threading.Thread(
             target=__upload_images_to_aws_thread,
             args=(
-                res, images_to_upload, project, annotation_status, prefix, thread_id,
-                chunksize, couldnt_upload, uploaded, tried_upload,
+                res, images_to_upload, project, annotation_status, prefix,
+                thread_id, chunksize, couldnt_upload, uploaded, tried_upload,
                 image_quality_in_editor, from_s3_bucket, folder_id
             ),
             daemon=True
@@ -987,9 +988,7 @@ def attach_image_urls_to_project(
     """
 
     project, folder = get_project_and_folder_metadata(project)
-    folder_name = project["name"] + (
-        f'/{folder["name"]}' if folder else ""
-    )
+    folder_name = project["name"] + (f'/{folder["name"]}' if folder else "")
     upload_state = common.upload_state_int_to_str(project.get("upload_state"))
     if upload_state == "Basic":
         raise SABaseException(
@@ -1027,8 +1026,7 @@ def attach_image_urls_to_project(
     image_data = pd.DataFrame(image_data, columns=["name", "url"])
     img_names_urls = image_data.values.tolist()
     logger.info(
-        "Uploading %s images to project %s.", len(img_names_urls),
-        folder_name
+        "Uploading %s images to project %s.", len(img_names_urls), folder_name
     )
     if len(img_names_urls) == 0:
         return ([], [], duplicate_images)
@@ -1038,14 +1036,14 @@ def attach_image_urls_to_project(
     else:
         folder_id = get_project_root_folder_id(project)
 
-    params = {'team_id': team_id , 'folder_id' : folder_id }
+    params = {'team_id': team_id, 'folder_id': folder_id}
     uploaded = [[] for _ in range(_NUM_THREADS)]
     tried_upload = [[] for _ in range(_NUM_THREADS)]
     couldnt_upload = [[] for _ in range(_NUM_THREADS)]
     finish_event = threading.Event()
 
-    res = _get_upload_auth_token(params=params,project_id=project_id)
-    
+    res = _get_upload_auth_token(params=params, project_id=project_id)
+
     prefix = res['filePath']
     limit = res['availableImageCount']
     images_to_upload = img_names_urls[:limit]
@@ -1084,7 +1082,7 @@ def attach_image_urls_to_project(
         for f in upload_thread:
             list_of_uploaded.append(str(f))
 
-    list_of_not_uploaded += [i[0] for i in images_to_skip ]
+    list_of_not_uploaded += [i[0] for i in images_to_skip]
     return (list_of_uploaded, list_of_not_uploaded, duplicate_images)