added --primary-monitor-only argument

This commit is contained in:
Koen van Eijk 2024-06-10 08:19:30 +02:00
parent 0b5a9ca124
commit ebf364f0c4
4 changed files with 23 additions and 2 deletions

View File

@ -82,6 +82,7 @@ Open your browser to:
## Arguments ## Arguments
`--storage-path` (default: user data path for your OS): allows you to specify the pathwhere the screenshots and database should be stored. `--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 ## Contribute

View File

@ -138,7 +138,7 @@ def search():
q = request.args.get("q") q = request.args.get("q")
entries = get_all_entries() entries = get_all_entries()
embeddings = [ 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) query_embedding = get_embedding(q)
similarities = [cosine_similarity(query_embedding, emb) for emb in embeddings] similarities = [cosine_similarity(query_embedding, emb) for emb in embeddings]

View File

@ -13,6 +13,14 @@ parser.add_argument(
help="Path to store the screenshots and database", 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() args = parser.parse_args()
if args.storage_path: if args.storage_path:

View File

@ -5,7 +5,7 @@ import mss
import numpy as np import numpy as np
from PIL import Image 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.database import insert_entry
from openrecall.nlp import get_embedding from openrecall.nlp import get_embedding
from openrecall.ocr import extract_text_from_image 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): def take_screenshots(monitor=1):
screenshots = [] screenshots = []
with mss.mss() as sct: with mss.mss() as sct:
for monitor in range(len(sct.monitors)): for monitor in range(len(sct.monitors)):
if args.primary_monitor_only and monitor != 1:
continue
monitor_ = sct.monitors[monitor] monitor_ = sct.monitors[monitor]
screenshot = np.array(sct.grab(monitor_)) screenshot = np.array(sct.grab(monitor_))
screenshot = screenshot[:, :, [2, 1, 0]] screenshot = screenshot[:, :, [2, 1, 0]]
screenshots.append(screenshot) screenshots.append(screenshot)
return screenshots return screenshots
def record_screenshots_thread(): def record_screenshots_thread():
last_screenshots = take_screenshots() last_screenshots = take_screenshots()
while True: while True:
screenshots = take_screenshots() screenshots = take_screenshots()
for i, screenshot in enumerate(screenshots): for i, screenshot in enumerate(screenshots):
last_screenshot = last_screenshots[i] last_screenshot = last_screenshots[i]
if not is_similar(screenshot, last_screenshot): if not is_similar(screenshot, last_screenshot):
last_screenshots[i] = screenshot last_screenshots[i] = screenshot
image = Image.fromarray(screenshot) image = Image.fromarray(screenshot)
@ -70,4 +81,5 @@ def record_screenshots_thread():
insert_entry( insert_entry(
text, timestamp, embedding, active_app_name, active_window_title text, timestamp, embedding, active_app_name, active_window_title
) )
time.sleep(3) time.sleep(3)