Merge pull request #52 from apirrone/dev

Merge dev into main

memento/background.py

@@ -19,9 +19,7 @@
 
 class Background:
     def __init__(self):
-        self.cache_path = os.path.join(os.environ["HOME"], ".cache", "memento")
-
-        if os.path.exists(os.path.join(self.cache_path, "0.json")):
+        if os.path.exists(os.path.join(utils.CACHE_PATH, "0.json")):
             print("EXISTING MEMENTO CACHE FOUND")
             print("Continue this recording or erase and start over ? ")
             print("1. Continue")
@@ -33,30 +31,30 @@ def __init__(self):
 
             if choice == "1":
                 self.nb_rec = len(
-                    [f for f in os.listdir(self.cache_path) if f.endswith(".mp4")]
+                    [f for f in os.listdir(utils.CACHE_PATH) if f.endswith(".mp4")]
                 )
                 self.frame_i = int(self.nb_rec * utils.FPS * utils.SECONDS_PER_REC)
             else:
-                os.system("rm -rf " + self.cache_path)
+                os.system("rm -rf " + utils.CACHE_PATH)
                 self.nb_rec = 0
                 self.frame_i = 0
         else:
             self.nb_rec = 0
             self.frame_i = 0
 
-        self.metadata_cache = MetadataCache(self.cache_path)
+        self.metadata_cache = MetadataCache()
 
-        os.makedirs(self.cache_path, exist_ok=True)
+        os.makedirs(utils.CACHE_PATH, exist_ok=True)
         self.db = Db()
         self.chromadb = Chroma(
-            persist_directory=self.cache_path,
+            persist_directory=utils.CACHE_PATH,
             embedding_function=OpenAIEmbeddings(),
             collection_name="memento_db",
         )
 
         self.sct = mss.mss()
         self.rec = utils.Recorder(
-            os.path.join(self.cache_path, str(self.nb_rec) + ".mp4")
+            os.path.join(utils.CACHE_PATH, str(self.nb_rec) + ".mp4")
         )
         self.rec.start()
 
@@ -221,5 +219,5 @@ def run(self):
                 self.rec.stop()
                 self.nb_rec += 1
                 self.rec = utils.Recorder(
-                    os.path.join(self.cache_path, str(self.nb_rec) + ".mp4")
+                    os.path.join(utils.CACHE_PATH, str(self.nb_rec) + ".mp4")
                 )

memento/caching.py

@@ -1,5 +1,5 @@
 import av
-from memento.utils import FPS, SECONDS_PER_REC, FRAME_CACHE_SIZE
+from memento.utils import FPS, SECONDS_PER_REC, FRAME_CACHE_SIZE, CACHE_PATH
 import time
 import os
 import json
@@ -22,10 +22,9 @@ def get_frame(self, frame_i):
 
 
 class ReadersCache:
-    def __init__(self, cache_path):
+    def __init__(self):
         self.readers = {}
         self.readers_ids = []  # in order to know the oldest reader
-        self.cache_path = cache_path
         self.cache_size = FRAME_CACHE_SIZE
 
     def select_video_id(self, frame_id):
@@ -37,7 +36,7 @@ def get_reader(self, frame_id):
         if video_id not in self.readers:  # Caching reader
             start = time.time()
             self.readers[video_id] = Reader(
-                os.path.join(self.cache_path, str(video_id) + ".mp4"), offset=offset
+                os.path.join(CACHE_PATH, str(video_id) + ".mp4"), offset=offset
             )
             self.readers_ids.append(video_id)
             # print(
@@ -78,8 +77,7 @@ def write(self, frame_id, data):
 
 
 class MetadataCache:
-    def __init__(self, cache_path):
-        self.cache_path = cache_path
+    def __init__(self):
         self.cache = {}
         self.cache_size = FRAME_CACHE_SIZE
         self.cache_ids = []
@@ -92,7 +90,7 @@ def get_metadata(self, frame_id):
         if metadata_id not in self.cache:
             start = time.time()
             self.cache[metadata_id] = Metadata(
-                os.path.join(self.cache_path, str(metadata_id) + ".json")
+                os.path.join(CACHE_PATH, str(metadata_id) + ".json")
             )
             # print(
             #     "Caching metadata",
@@ -112,7 +110,7 @@ def get_metadata(self, frame_id):
     def get_frame_metadata(self, frame_id):
         metadata = self.get_metadata(frame_id)
         return metadata.get_frame(frame_id)
-    
+
     def write(self, frame_id, data):
         metadata = self.get_metadata(frame_id)
         metadata.write(frame_id, data)

memento/db.py

@@ -1,11 +1,11 @@
 import sqlite3
 import os
+from memento.utils import CACHE_PATH
 
 
 class Db:
     def __init__(self):
-        self.cache_path = os.path.join(os.environ["HOME"], ".cache", "memento")
-        db_path = os.path.join(self.cache_path, "memento.db")
+        db_path = os.path.join(CACHE_PATH, "memento.db")
         create_tables = False
         if not os.path.isfile(db_path):
             create_tables = True
@@ -70,7 +70,7 @@ def search(self, query):
 
         results = {}
         for row in cursor:
-            print(row)
+            # print(row)
             rank = row[0]
             frame_id = str(row[1])
             if frame_id not in results:

memento/timeline/apps.py

@@ -0,0 +1,60 @@
+import numpy as np
+from memento.timeline.icon_getter import IconGetter
+import time
+
+COLOR_PALETTE = [
+    [244, 67, 54],
+    [233, 30, 99],
+    [156, 39, 176],
+    [103, 58, 183],
+    [63, 81, 181],
+    [33, 150, 243],
+    [3, 169, 244],
+    [0, 188, 212],
+    [0, 150, 136],
+    [76, 175, 80],
+    [139, 195, 74],
+    [205, 220, 57],
+    [255, 235, 59],
+    [255, 193, 7],
+    [255, 152, 0],
+    [255, 87, 34],
+]
+# TODO add more colors if needed ?
+
+
+class Apps:
+    def __init__(self, frame_getter):
+        self.apps = {}
+        self.frame_getter = frame_getter
+        ws = self.frame_getter.window_size
+        self.h = ws[1] // 40
+
+        self.metadata_cache = self.frame_getter.metadata_cache
+        self.nb_frames = self.frame_getter.nb_frames
+        self.ig = IconGetter(size=self.h)
+
+        # Hacky solution for faster timeline startup
+        stride = 100
+        if len(self.ig.icon_cache) == 0 or self.nb_frames < 1000:
+            stride = 1
+        for i in range(0, self.nb_frames, stride):
+            app = self.metadata_cache.get_frame_metadata(i)["window_title"]
+            if app not in self.apps:
+                self.apps[app] = {}
+                if len(self.apps) < len(COLOR_PALETTE):
+                    self.apps[app]["color"] = tuple(COLOR_PALETTE[len(self.apps)])
+                else:
+                    self.apps[app]["color"] = tuple(np.random.randint(0, 255, size=3))
+                icon_small, icon_big = self.ig.lookup_icon(app)
+                self.apps[app]["icon_small"] = icon_small
+                self.apps[app]["icon_big"] = icon_big
+
+    def get_color(self, app):
+        return self.apps[app]["color"]
+
+    def get_icon(self, app, small=True):
+        if small:
+            return self.apps[app]["icon_small"]
+        else:
+            return self.apps[app]["icon_big"]

memento/timeline/chat.py

@@ -18,9 +18,13 @@
 # Chat window on the right of the screen
 class Chat:
     def __init__(self, frame_getter):
-        self.cache_path = os.path.join(os.environ["HOME"], ".cache", "memento")
-
         self.frame_getter = frame_getter
+        self.key_ok = True
+        self.msg = ""
+        if "OPENAI_API_KEY" not in os.environ:
+            self.key_ok = False
+            self.msg = "Chat requires an OpenAI API key to work. Set the OPENAI_API_KEY env variable to your API key."
+            print(self.msg)
 
         ws = self.frame_getter.window_size
         self.w = ws[0] // 3
@@ -48,7 +52,7 @@ def __init__(self, frame_getter):
         self.frame_peek_hovered_id = None
 
         self.chromadb = Chroma(
-            persist_directory=self.cache_path,
+            persist_directory=utils.CACHE_PATH,
             embedding_function=OpenAIEmbeddings(),
             collection_name="memento_db",
         )
@@ -61,7 +65,7 @@ def __init__(self, frame_getter):
         # Define prompt
         template = """Use the following pieces of context and metadata to answer the question at the end. Answer in the same language the question was asked.
         If you don't know the answer, just say that you don't know, don't try to make up an answer.
-        You will format your answer in json, with the keys "answer" and "frames_ids". 
+        You will format your answer in json, with the keys "answer" and "frames_ids". Always include these keys, even if you did not find anything. Just say it and return an empty list for "frames_ids".
         The value of "answer" will be the answer to the question, and the value of "frames_ids" will be a list of frame_ids from which you got the information from using the metadata.
         Use three sentences maximum and keep the answer as concise as possible.
 
@@ -76,7 +80,7 @@ def __init__(self, frame_getter):
         )
 
         self.qa = ConversationalRetrievalChain.from_llm(
-            ChatOpenAI(model_name="gpt-4", temperature=0.8),
+            ChatOpenAI(model_name="gpt-3.5-turbo-0301", temperature=0.8),
             self.retriever,
             memory=self.memory,
             verbose=True,
@@ -122,14 +126,22 @@ def process_chat_query(self):
                 }
 
             result = self.qa(inputs={"question": inp, "md": md})
-            result = json.loads(result["answer"])
+
+            try:
+                result = json.loads(result["answer"])
+            except json.decoder.JSONDecodeError as e:
+                print("Error decoding json:", e)
+                result = {"answer": "Error decoding json", "frames_ids": []}
+
             print("Answer:", result["answer"])
             print("frames_ids:", result["frames_ids"])
             self.answer_queue.put(result)
 
     def events(self, events):
         if not self.active:
             return None
+        if not self.key_ok:
+            return None
 
         self.textinput.update(events)
         for event in events:
@@ -344,5 +356,7 @@ def draw(self, screen):
         screen.blit(surf, (self.x, self.y))
 
         self.draw_chat_history(screen)
+        if not self.key_ok:
+            self.draw_bubble(screen, self.msg, 0, question=False)
         self.handle_frames_peeks(screen)
         self.draw_input_box(screen)

memento/timeline/frame_getter.py

@@ -10,14 +10,13 @@ class FrameGetter:
     def __init__(self, window_size):
         self.window_size = window_size
 
-        self.cache_path = os.path.join(os.environ["HOME"], ".cache", "memento")
-        self.readers_cache = ReadersCache(self.cache_path)
-        self.metadata_cache = MetadataCache(self.cache_path)
+        self.readers_cache = ReadersCache()
+        self.metadata_cache = MetadataCache()
         self.annotations = {}
         self.current_ret_annotated = 0
         self.nb_frames = int(
             (
-                len(glob(os.path.join(self.cache_path, "*.mp4")))
+                len(glob(os.path.join(utils.CACHE_PATH, "*.mp4")))
                 * utils.FPS
                 * utils.SECONDS_PER_REC
             )
@@ -32,9 +31,13 @@ def toggle_debug_mode(self):
         self.debug_mode = not self.debug_mode
         self.clear_annotations()
 
-    def get_frame(self, frame_i):
+    def get_frame(self, frame_i, resize=None):
         im = self.current_displayed_frame
 
+        # Resize frame if needed, still use cache
+        if im is not None and resize != im.shape:
+            self.current_displayed_frame = None
+
         # Avoid resizing and converting the same frame each time
         if (
             frame_i != self.current_displayed_frame_i
@@ -43,7 +46,10 @@ def get_frame(self, frame_i):
             im = self.readers_cache.get_frame(min(self.nb_frames - 1, frame_i))
             self.process_debug(frame_i)
             im = self.annotate_frame(frame_i, im)
-            im = cv2.resize(im, self.window_size)
+            if resize:
+                im = cv2.resize(im, resize)
+            else:
+                im = cv2.resize(im, self.window_size)
             im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB).swapaxes(0, 1)
             self.current_displayed_frame = im
             self.current_displayed_frame_i = frame_i
@@ -130,10 +136,20 @@ def get_next_annotated_frame_i(self):
         else:
             return 0
 
-    def set_annotation(self, annotations):
+    def set_annotations(self, annotations):
         self.annotations = annotations
         self.current_displayed_frame = None
 
+    def get_annotations(self):
+        return self.annotations
+
+    def get_annotated_frames(self):
+        frames = []
+        for frame_i in list(self.annotations.keys())[:10]:
+            frames.append(self.get_frame(int(frame_i)))
+
+        return frames
+
     def get_annotations_text(self):
         text = ""
         for entries in self.annotations.values():

memento/timeline/icon_getter.py

@@ -2,6 +2,8 @@
 from glob import glob
 import os
 import pygame
+from memento.utils import CACHE_PATH
+import pickle
 
 
 class IconGetter:
@@ -10,10 +12,33 @@ def __init__(self, size=100):
         self.icons_paths = ["/usr/share/icons/**/", "/usr/share/pixmaps/**/"]
         self.extensions = [".png", ".jpg"]
         self.size = size
+        self.icon_cache_path = os.path.join(CACHE_PATH, "icon_cache.pickle")
+        if os.path.exists(self.icon_cache_path):
+            self.icon_cache = pickle.load(open(self.icon_cache_path, "rb"))
+        else:
+            self.icon_cache = {}
 
     def lookup_icon(self, app_name):
-        if app_name in [None, "None"]:
+        if app_name not in self.icon_cache:
+            best_icon_path = self.new_icon(app_name)
+        else:
+            best_icon_path = self.icon_cache[app_name]
+
+        if best_icon_path is None:
             return None, None
+
+        icon_big = pygame.transform.scale(
+            pygame.image.load(best_icon_path), (self.size * 2, self.size * 2)
+        )
+        icon_small = pygame.transform.scale(
+            pygame.image.load(best_icon_path), (self.size, self.size)
+        )
+
+        return icon_small, icon_big
+
+    def new_icon(self, app_name):
+        if app_name in [None, "None"]:
+            return None
         if self.all_icons_paths == []:
             for icons_path in self.icons_paths:
                 for ext in self.extensions:
@@ -28,14 +53,7 @@ def lookup_icon(self, app_name):
                 best_score = score
                 best_icon_path = icon_path
 
-        if best_icon_path is None:
-            return None, None
-
-        icon_big = pygame.transform.scale(
-            pygame.image.load(best_icon_path), (self.size * 2, self.size * 2)
-        )
-        icon_small = pygame.transform.scale(
-            pygame.image.load(best_icon_path), (self.size, self.size)
-        )
+        self.icon_cache[app_name] = best_icon_path
+        pickle.dump(self.icon_cache, open(self.icon_cache_path, "wb"))
 
-        return icon_small, icon_big
+        return best_icon_path