From 3a3b788c6542f363a1ebe5486343bfd10efc80a9 Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Sat, 19 Jun 2021 08:38:42 -0500
Subject: [PATCH] save clean snapshot

---
 frigate/object_processing.py | 38 ++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/frigate/object_processing.py b/frigate/object_processing.py
index b483940f3..4968f14f1 100644
--- a/frigate/object_processing.py
+++ b/frigate/object_processing.py
@@ -192,6 +192,27 @@ class TrackedObject:
             ret, jpg = cv2.imencode(".jpg", np.zeros((175, 175, 3), np.uint8))
             return jpg.tobytes()
 
+    def get_clean_png(self):
+        if self.thumbnail_data is None:
+            return None
+
+        try:
+            best_frame = cv2.cvtColor(
+                self.frame_cache[self.thumbnail_data["frame_time"]],
+                cv2.COLOR_YUV2BGR_I420,
+            )
+        except KeyError:
+            logger.warning(
+                f"Unable to create clean png because frame {self.thumbnail_data['frame_time']} is not in the cache"
+            )
+            return None
+
+        ret, png = cv2.imencode(".png", best_frame)
+        if ret:
+            return png.tobytes()
+        else:
+            return None
+
     def get_jpg_bytes(
         self, timestamp=False, bounding_box=False, crop=False, height=None
     ):
@@ -615,6 +636,23 @@ class TrackedObjectProcessor(threading.Thread):
                         ) as j:
                             j.write(jpg_bytes)
                         event_data["has_snapshot"] = True
+
+                    # write clean snapshot if enabled
+                    if snapshot_config.clean_copy:
+                        png_bytes = obj.get_clean_png()
+                        if png_bytes is None:
+                            logger.warning(
+                                f"Unable to save clean snapshot for {obj.obj_data['id']}."
+                            )
+                        else:
+                            with open(
+                                os.path.join(
+                                    CLIPS_DIR,
+                                    f"{camera}-{obj.obj_data['id']}-clean.png",
+                                ),
+                                "wb",
+                            ) as p:
+                                p.write(png_bytes)
             self.event_queue.put(("end", camera, event_data))
 
         def snapshot(camera, obj: TrackedObject, current_frame_time):