mirror of
https://github.com/tcsenpai/pensieve.git
synced 2025-06-06 03:05:25 +00:00
feat: use mss instead of imagegrab
This commit is contained in:
parent
c24d9cc2e9
commit
30f9a45d8a
@ -42,6 +42,7 @@ dependencies = [
|
||||
"timm",
|
||||
"einops",
|
||||
"modelscope",
|
||||
"mss",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
|
@ -4,11 +4,11 @@ import time
|
||||
import logging
|
||||
import platform
|
||||
import subprocess
|
||||
from PIL import Image, ImageGrab
|
||||
from PIL import Image
|
||||
import imagehash
|
||||
from memos.utils import write_image_metadata
|
||||
from screeninfo import get_monitors
|
||||
import ctypes
|
||||
from mss import mss
|
||||
|
||||
if platform.system() == "Windows":
|
||||
import win32gui
|
||||
@ -173,57 +173,49 @@ def take_screenshot_windows(
|
||||
app_name,
|
||||
window_title,
|
||||
):
|
||||
for monitor in get_monitors():
|
||||
safe_monitor_name = "".join(
|
||||
c for c in monitor.name if c.isalnum() or c in ("_", "-")
|
||||
)
|
||||
logging.info(f"Processing monitor: {safe_monitor_name}")
|
||||
with mss() as sct:
|
||||
for i, monitor in enumerate(sct.monitors[1:], 1): # Skip the first monitor (entire screen)
|
||||
safe_monitor_name = f"monitor_{i}"
|
||||
logging.info(f"Processing monitor: {safe_monitor_name}")
|
||||
|
||||
webp_filename = os.path.join(
|
||||
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,
|
||||
webp_filename = os.path.join(
|
||||
base_dir, date, f"screenshot-{timestamp}-of-{safe_monitor_name}.webp"
|
||||
)
|
||||
)
|
||||
img = img.convert("RGB")
|
||||
current_hash = str(imagehash.phash(img))
|
||||
|
||||
if (
|
||||
safe_monitor_name in previous_hashes
|
||||
and imagehash.hex_to_hash(current_hash)
|
||||
- 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."
|
||||
img = sct.grab(monitor)
|
||||
img = Image.frombytes("RGB", img.size, img.bgra, "raw", "BGRX")
|
||||
current_hash = str(imagehash.phash(img))
|
||||
|
||||
if (
|
||||
safe_monitor_name in previous_hashes
|
||||
and imagehash.hex_to_hash(current_hash)
|
||||
- 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
|
||||
screen_sequences[safe_monitor_name] = (
|
||||
screen_sequences.get(safe_monitor_name, 0) + 1
|
||||
)
|
||||
metadata = {
|
||||
"timestamp": timestamp,
|
||||
"active_app": app_name,
|
||||
"active_window": window_title,
|
||||
"screen_name": safe_monitor_name,
|
||||
"sequence": screen_sequences[safe_monitor_name],
|
||||
}
|
||||
|
||||
metadata = {
|
||||
"timestamp": timestamp,
|
||||
"active_app": app_name,
|
||||
"active_window": window_title,
|
||||
"screen_name": safe_monitor_name,
|
||||
"sequence": screen_sequences[safe_monitor_name],
|
||||
}
|
||||
img.save(webp_filename, format="WebP", quality=85)
|
||||
write_image_metadata(webp_filename, metadata)
|
||||
save_screen_sequences(base_dir, screen_sequences, date)
|
||||
|
||||
img.save(webp_filename, format="WebP", quality=85)
|
||||
write_image_metadata(webp_filename, metadata)
|
||||
save_screen_sequences(base_dir, screen_sequences, date)
|
||||
|
||||
yield safe_monitor_name, webp_filename, "Saved"
|
||||
yield safe_monitor_name, webp_filename, "Saved"
|
||||
|
||||
|
||||
def take_screenshot(
|
||||
|
Loading…
x
Reference in New Issue
Block a user