From ebf364f0c457558db27276589c4bc53cc1979a8d Mon Sep 17 00:00:00 2001
From: Koen van Eijk <vaneijk.koen@gmail.com>
Date: Mon, 10 Jun 2024 08:19:30 +0200
Subject: [PATCH] added --primary-monitor-only argument

---
 README.md                |  1 +
 openrecall/app.py        |  2 +-
 openrecall/config.py     |  8 ++++++++
 openrecall/screenshot.py | 14 +++++++++++++-
 4 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 277f546..ed8789f 100644
--- a/README.md
+++ b/README.md
@@ -82,6 +82,7 @@ Open your browser to:
 
 ## Arguments
 `--storage-path` (default: user data path for your OS): allows you to specify the pathwhere the screenshots and database should be stored.
+'--primary-monitor-only' (default: False): only record the primary monitor (rather than individual screenshots for other monitors)
 
 ## Contribute
 
diff --git a/openrecall/app.py b/openrecall/app.py
index b13fe73..6ed35bc 100644
--- a/openrecall/app.py
+++ b/openrecall/app.py
@@ -138,7 +138,7 @@ def search():
     q = request.args.get("q")
     entries = get_all_entries()
     embeddings = [
-        np.frombuffer(entry["embedding"], dtype=np.float64) for entry in entries
+        np.frombuffer(entry.embedding, dtype=np.float64) for entry in entries
     ]
     query_embedding = get_embedding(q)
     similarities = [cosine_similarity(query_embedding, emb) for emb in embeddings]
diff --git a/openrecall/config.py b/openrecall/config.py
index eb05d4a..0b67ab0 100644
--- a/openrecall/config.py
+++ b/openrecall/config.py
@@ -13,6 +13,14 @@ parser.add_argument(
     help="Path to store the screenshots and database",
 )
 
+parser.add_argument(
+    "--primary-monitor-only",
+    action="store_true",
+    help="Only record the primary monitor",
+    type=bool,
+    default=False,
+)
+
 args = parser.parse_args()
 
 if args.storage_path:
diff --git a/openrecall/screenshot.py b/openrecall/screenshot.py
index e0ced3e..02bbadf 100644
--- a/openrecall/screenshot.py
+++ b/openrecall/screenshot.py
@@ -5,7 +5,7 @@ import mss
 import numpy as np
 from PIL import Image
 
-from openrecall.config import screenshots_path
+from openrecall.config import screenshots_path, args
 from openrecall.database import insert_entry
 from openrecall.nlp import get_embedding
 from openrecall.ocr import extract_text_from_image
@@ -39,21 +39,32 @@ def is_similar(img1, img2, similarity_threshold=0.9):
 
 def take_screenshots(monitor=1):
     screenshots = []
+
     with mss.mss() as sct:
         for monitor in range(len(sct.monitors)):
+
+            if args.primary_monitor_only and monitor != 1:
+                continue
+            
             monitor_ = sct.monitors[monitor]
             screenshot = np.array(sct.grab(monitor_))
             screenshot = screenshot[:, :, [2, 1, 0]]
             screenshots.append(screenshot)
+    
     return screenshots
 
 
 def record_screenshots_thread():
     last_screenshots = take_screenshots()
+
     while True:
+    
         screenshots = take_screenshots()
+    
         for i, screenshot in enumerate(screenshots):
+            
             last_screenshot = last_screenshots[i]
+
             if not is_similar(screenshot, last_screenshot):
                 last_screenshots[i] = screenshot
                 image = Image.fromarray(screenshot)
@@ -70,4 +81,5 @@ def record_screenshots_thread():
                 insert_entry(
                     text, timestamp, embedding, active_app_name, active_window_title
                 )
+                
         time.sleep(3)