tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-07 03:35:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: use mss instead of imagegrab

Browse Source
This commit is contained in:
arkohut 2024-09-11 17:50:33 +08:00
parent c24d9cc2e9
commit 30f9a45d8a
2 changed files with 39 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
pyproject.toml
View File

@ -42,6 +42,7 @@ dependencies = [
"timm", "timm",
"einops", "einops",
"modelscope", "modelscope",
"mss",
] ]
[project.urls] [project.urls]

84
screen_recorder/common.py
View File

@ -4,11 +4,11 @@ import time
import logging import logging
import platform import platform
import subprocess import subprocess
from PIL import Image, ImageGrab from PIL import Image
import imagehash import imagehash
from memos.utils import write_image_metadata from memos.utils import write_image_metadata
from screeninfo import get_monitors
import ctypes import ctypes
from mss import mss
if platform.system() == "Windows": if platform.system() == "Windows":
import win32gui import win32gui
@ -173,57 +173,49 @@ def take_screenshot_windows(
app_name, app_name,
window_title, window_title,
): ):
for monitor in get_monitors(): with mss() as sct:
safe_monitor_name = "".join( for i, monitor in enumerate(sct.monitors[1:], 1): # Skip the first monitor (entire screen)
c for c in monitor.name if c.isalnum() or c in ("_", "-") safe_monitor_name = f"monitor_{i}"
) logging.info(f"Processing monitor: {safe_monitor_name}")
logging.info(f"Processing monitor: {safe_monitor_name}")
webp_filename = os.path.join( webp_filename = os.path.join(
base_dir, date, f"screenshot-{timestamp}-of-{safe_monitor_name}.webp" base_dir, date, f"screenshot-{timestamp}-of-{safe_monitor_name}.webp"
)
img = ImageGrab.grab(
bbox=(
monitor.x,
monitor.y,
monitor.x + monitor.width,
monitor.y + monitor.height,
) )
)
img = img.convert("RGB")
current_hash = str(imagehash.phash(img))
if ( img = sct.grab(monitor)
safe_monitor_name in previous_hashes img = Image.frombytes("RGB", img.size, img.bgra, "raw", "BGRX")
and imagehash.hex_to_hash(current_hash) current_hash = str(imagehash.phash(img))
- imagehash.hex_to_hash(previous_hashes[safe_monitor_name])
< threshold if (
): safe_monitor_name in previous_hashes
logging.info( and imagehash.hex_to_hash(current_hash)
f"Screenshot for {safe_monitor_name} is similar to the previous one. Skipping." - imagehash.hex_to_hash(previous_hashes[safe_monitor_name])
< threshold
):
logging.info(
f"Screenshot for {safe_monitor_name} is similar to the previous one. Skipping."
)
yield safe_monitor_name, None, "Skipped (similar to previous)"
continue
previous_hashes[safe_monitor_name] = current_hash
screen_sequences[safe_monitor_name] = (
screen_sequences.get(safe_monitor_name, 0) + 1
) )
yield safe_monitor_name, None, "Skipped (similar to previous)"
continue
previous_hashes[safe_monitor_name] = current_hash metadata = {
screen_sequences[safe_monitor_name] = ( "timestamp": timestamp,
screen_sequences.get(safe_monitor_name, 0) + 1 "active_app": app_name,
) "active_window": window_title,
"screen_name": safe_monitor_name,
"sequence": screen_sequences[safe_monitor_name],
}
metadata = { img.save(webp_filename, format="WebP", quality=85)
"timestamp": timestamp, write_image_metadata(webp_filename, metadata)
"active_app": app_name, save_screen_sequences(base_dir, screen_sequences, date)
"active_window": window_title,
"screen_name": safe_monitor_name,
"sequence": screen_sequences[safe_monitor_name],
}
img.save(webp_filename, format="WebP", quality=85) yield safe_monitor_name, webp_filename, "Saved"
write_image_metadata(webp_filename, metadata)
save_screen_sequences(base_dir, screen_sequences, date)
yield safe_monitor_name, webp_filename, "Saved"
def take_screenshot( def take_screenshot(