enhancing Frame implementation

SSusantAchary · Nov 12, 2020 · 0b1a45b · 0b1a45b
1 parent 9ca900f
commit 0b1a45b
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 33 deletions.
diff --git a/fiftyone/core/frame.py b/fiftyone/core/frame.py
@@ -16,7 +16,13 @@
 
 import fiftyone as fo
 from fiftyone.core.document import Document
+import fiftyone.core.fields as fof
 import fiftyone.core.frame_utils as fofu
+import fiftyone.core.labels as fol
+from fiftyone.core.odm.document import (
+    DynamicEmbeddedDocument,
+    EmbeddedDocument,
+)
 from fiftyone.core.odm.frame import (
     NoDatasetFrameSampleDocument,
     DatasetFrameSampleDocument,
@@ -45,7 +51,6 @@ def __init__(self):
         self._iter = None
         self._iter_frame = None
         self._replacements = {}
-        self._replacements_complete = False
 
     def __str__(self):
         return "<%s: %s>" % (self.__class__.__name__, fou.pformat(dict(self)))
@@ -56,7 +61,7 @@ def __repr__(self):
         return "{ <%d frame%s> }" % (num_frames, plural)
 
     def __len__(self):
-        return self._sample._doc.frames["frame_count"]
+        return self._sample._doc.frames.frame_count
 
     def _set_replacement(self, frame):
         self._replacements[frame.frame_number] = frame
@@ -125,7 +130,7 @@ def __contains__(self, frame_number):
         if frame_number in self._replacements:
             return True
 
-        if self._sample._in_db:
+        if self._in_db:
             find_d = self._make_filter(frame_number, self._sample._id)
             return self._frame_collection.find_one(find_d) is not None
 
@@ -134,7 +139,6 @@ def __contains__(self, frame_number):
     def __getitem__(self, frame_number):
         fofu.validate_frame_number(frame_number)
 
-        dataset = None
         frame = None
         d = None
         default_d = {
@@ -148,11 +152,11 @@ def __getitem__(self, frame_number):
             frame = self._iter_frame
         elif frame_number in self._replacements:
             frame = self._replacements[frame_number]
-        elif self._sample._in_db:
+        elif self._in_db:
             d = self._frame_collection.find_one(default_d)
             if d is None:
                 d = default_d
-                self._sample._doc.frames["frame_count"] += 1
+                self._sample._doc.frames.frame_count += 1
 
             frame = Frame.from_doc(
                 self._sample._dataset._frame_dict_to_doc(d),
@@ -161,12 +165,9 @@ def __getitem__(self, frame_number):
             self._set_replacement(frame)
         else:
             frame = Frame.from_doc(NoDatasetFrameSampleDocument(**default_d))
-            self._sample._doc.frames["frame_count"] += 1
+            self._sample._doc.frames.frame_count += 1
             self._set_replacement(frame)
 
-        if self._sample._in_db:
-            dataset = self._sample._dataset
-
         return frame
 
     def __setitem__(self, frame_number, frame):
@@ -179,9 +180,9 @@ def __setitem__(self, frame_number, frame):
         doc.set_field("frame_number", frame_number)
         doc._sample_id = self._sample._id
 
-        if not self._sample._in_db and frame_number not in self._replacements:
-            self._sample._doc.frames["frame_count"] += 1
-        elif self._sample._in_db:
+        if not self._in_db and frame_number not in self._replacements:
+            self._sample._doc.frames.frame_count += 1
+        elif self._in_db:
             if (
                 self._iter_frame is not None
                 and frame_number != self._iter_frame.frame_number
@@ -192,7 +193,7 @@ def __setitem__(self, frame_number, frame):
                 }
                 exists = self._frame_collection.find(find_d)
                 if exists is None:
-                    self._sample._doc.frames["frame_count"] += 1
+                    self._sample._doc.frames.frame_count += 1
 
         self._set_replacement(frame)
 
@@ -287,27 +288,40 @@ def merge(self, frames, overwrite=True):
             else:
                 self[frame_number] = frame
 
-    @property
-    def _first_frame(self):
-        first_frame = self._replacements.get(1, None)
-        if first_frame is None and self._sample._in_db:
-            first_frame = self._sample._doc.frames["first_frame"]
-
-        return first_frame
+    def clear(self):
+        """Removes all frames from this instance."""
+        self._replacements = {}
+        if self._in_db:
+            self._frame_collection.delete_many(
+                {"_sample_id": self._sample._id}
+            )
+            self._sample._doc.clear_field("frames")
 
     def to_mongo_dict(self):
         first_frame = self._first_frame
         return {
-            "frame_count": self._sample._doc.frames["frame_count"],
+            "frame_count": self._sample._doc.frames.frame_count,
             "first_frame": self._first_frame,
         }
 
+    @property
+    def _in_db(self):
+        return self._sample._in_db
+
+    @property
+    def _first_frame(self):
+        first_frame = self._replacements.get(1, None)
+        if first_frame is None and self._in_db:
+            first_frame = self._sample._doc.frames.first_frame
+
+        return first_frame
+
     @property
     def _frame_collection(self):
         return self._sample._dataset._frame_collection
 
     def _iter_frames(self):
-        if self._sample._in_db:
+        if self._in_db:
             repl_fns = sorted(self._replacements.keys())
             repl_fn = repl_fns[0] if len(repl_fns) else None
             find_d = {"_sample_id": self._sample._id}
@@ -329,6 +343,12 @@ def _iter_frames(self):
                 self._set_replacement(self._iter_frame)
                 yield self._iter_frame
 
+    def _to_frames_dict(self):
+        return {
+            str(frame_number): frame.to_dict()
+            for frame_number, frame in self.items()
+        }
+
     def _get_field_cls(self):
         return self._sample._doc.frames.__class__
 
@@ -341,12 +361,9 @@ def _get_first_frame(self):
         return None
 
     def _save(self, insert=False):
-        if not self._sample._in_db:
+        if not self._in_db:
             return
 
-        # @todo avoid local import?
-        from fiftyone.core.labels import _FrameLabels
-
         d = self._get_first_frame()
         if d is not None:
             for k, v in d.items():
@@ -367,7 +384,7 @@ def _save(self, insert=False):
                 else:
                     d[k] = v
 
-            self._sample._doc.frames.first_frame = _FrameLabels(**d)
+            self._sample._doc.frames.first_frame = fol._FrameLabels(**d)
 
         self._save_replacements(insert)
 
@@ -415,6 +432,9 @@ def __setitem__(self, field_name, value):
     def _skip_iter_field_names(self):
         return ("frame_number",)
 
+    def _get_field_names(self, include_private=False):
+        return self._doc._get_field_names(include_private=include_private)
+
     def copy(self):
         """Returns a deep copy of the frame that has not been added to the
         database.
@@ -475,6 +495,12 @@ def _purge_field(cls, collection_name, field_name):
             for document in samples.values():
                 document._doc._data.pop(field_name, None)
 
+    @classmethod
+    def _reload_docs(cls, collection_name):
+        for samples in cls._instances[collection_name].values():
+            for document in samples.values():
+                document.reload()
+
     def _set_backing_doc(self, doc, dataset=None):
         """Sets the backing doc for the sample.
 

diff --git a/fiftyone/utils/data/exporters.py b/fiftyone/utils/data/exporters.py
@@ -880,11 +880,7 @@ def close(self, *args):
         )
 
     def _export_frame_labels(self, sample, uuid):
-        frames = {}
-        for frame_number, frame in sample.frames.items():
-            frames[str(frame_number)] = frame.to_dict()
-
-        frames_dict = {"frames": frames}
+        frames_dict = {"frames": sample.frames._to_frames_dict()}
         outpath = os.path.join(self._frame_labels_dir, uuid + ".json")
         etas.write_json(frames_dict, outpath, pretty_print=self.pretty_print)