From 5e84f4566ac97b3cb48d99f435b0e894457c3e8a Mon Sep 17 00:00:00 2001
From: Petr Kalis <petr.kalis@gmail.com>
Date: Tue, 8 Mar 2022 13:50:42 +0100
Subject: [PATCH] OP-2860 - added possibility to get number of frames from
 video file with ffprobe

Previously wrong hardcoded value was used. This implementation needs to be monitored for weird format of published video files.
---
 .../publish/collect_published_files.py        | 71 ++++++++++++++++---
 1 file changed, 63 insertions(+), 8 deletions(-)

diff --git a/openpype/hosts/webpublisher/plugins/publish/collect_published_files.py b/openpype/hosts/webpublisher/plugins/publish/collect_published_files.py
index abad14106f..8b21842635 100644
--- a/openpype/hosts/webpublisher/plugins/publish/collect_published_files.py
+++ b/openpype/hosts/webpublisher/plugins/publish/collect_published_files.py
@@ -10,14 +10,18 @@ Provides:
 import os
 import clique
 import tempfile
+import math
+
 from avalon import io
 import pyblish.api
-from openpype.lib import prepare_template_data
+from openpype.lib import prepare_template_data, get_asset, ffprobe_streams
+from openpype.lib.vendor_bin_utils import get_fps
 from openpype.lib.plugin_tools import (
     parse_json,
     get_subset_name_with_asset_doc
 )
 
+
 class CollectPublishedFiles(pyblish.api.ContextPlugin):
     """
     This collector will try to find json files in provided
@@ -49,10 +53,7 @@ class CollectPublishedFiles(pyblish.api.ContextPlugin):
         self.log.info("task_sub:: {}".format(task_subfolders))
 
         asset_name = context.data["asset"]
-        asset_doc = io.find_one({
-            "type": "asset",
-            "name": asset_name
-        })
+        asset_doc = get_asset()
         task_name = context.data["task"]
         task_type = context.data["taskType"]
         project_name = context.data["project_name"]
@@ -97,11 +98,26 @@ class CollectPublishedFiles(pyblish.api.ContextPlugin):
                 instance.data["frameEnd"] = \
                     instance.data["representations"][0]["frameEnd"]
             else:
-                instance.data["frameStart"] = 0
-                instance.data["frameEnd"] = 1
+                frame_start = asset_doc["data"]["frameStart"]
+                instance.data["frameStart"] = frame_start
+                instance.data["frameEnd"] = asset_doc["data"]["frameEnd"]
                 instance.data["representations"] = self._get_single_repre(
                     task_dir, task_data["files"], tags
                 )
+                file_url = os.path.join(task_dir, task_data["files"][0])
+                duration = self._get_duration(file_url)
+                if duration:
+                    try:
+                        frame_end = int(frame_start) + math.ceil(duration)
+                        instance.data["frameEnd"] = math.ceil(frame_end)
+                        self.log.debug("frameEnd:: {}".format(
+                            instance.data["frameEnd"]))
+                    except ValueError:
+                        self.log.warning("Unable to count frames "
+                                         "duration {}".format(duration))
+
+            instance.data["handleStart"] = asset_doc["data"]["handleStart"]
+            instance.data["handleEnd"] = asset_doc["data"]["handleEnd"]
 
             self.log.info("instance.data:: {}".format(instance.data))
 
@@ -127,7 +143,7 @@ class CollectPublishedFiles(pyblish.api.ContextPlugin):
         return [repre_data]
 
     def _process_sequence(self, files, task_dir, tags):
-        """Prepare reprentations for sequence of files."""
+        """Prepare representation for sequence of files."""
         collections, remainder = clique.assemble(files)
         assert len(collections) == 1, \
             "Too many collections in {}".format(files)
@@ -188,6 +204,7 @@ class CollectPublishedFiles(pyblish.api.ContextPlugin):
         msg = "No family found for combination of " +\
               "task_type: {}, is_sequence:{}, extension: {}".format(
                   task_type, is_sequence, extension)
+        found_family = "render"
         assert found_family, msg
 
         return (found_family,
@@ -243,3 +260,41 @@ class CollectPublishedFiles(pyblish.api.ContextPlugin):
             return version[0].get("version") or 0
         else:
             return 0
+
+    def _get_duration(self, file_url):
+        """Return duration in frames"""
+        try:
+            streams = ffprobe_streams(file_url, self.log)
+        except Exception as exc:
+            raise AssertionError((
+                "FFprobe couldn't read information about input file: \"{}\"."
+                " Error message: {}"
+            ).format(file_url, str(exc)))
+
+        first_video_stream = None
+        for stream in streams:
+            if "width" in stream and "height" in stream:
+                first_video_stream = stream
+                break
+
+        if first_video_stream:
+            nb_frames = stream.get("nb_frames")
+            if nb_frames:
+                try:
+                    return int(nb_frames)
+                except ValueError:
+                    self.log.warning(
+                        "nb_frames {} not convertible".format(nb_frames))
+
+                    duration = stream.get("duration")
+                    frame_rate = get_fps(stream.get("r_frame_rate", '0/0'))
+                    self.log.debu("duration:: {} frame_rate:: {}".format(
+                        duration, frame_rate))
+                    try:
+                        return float(duration) * float(frame_rate)
+                    except ValueError:
+                        self.log.warning(
+                            "{} or {} cannot be converted".format(duration,
+                                                                  frame_rate))
+
+        self.log.warning("Cannot get number of frames")