From ad8eb637230ac04fa89402756ef41b3ca8aa71f6 Mon Sep 17 00:00:00 2001 From: apirrone Date: Sat, 9 Sep 2023 20:30:54 +0200 Subject: [PATCH 1/7] reducing bitrate --- memento/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/memento/utils.py b/memento/utils.py index d18b452..825914f 100644 --- a/memento/utils.py +++ b/memento/utils.py @@ -64,7 +64,7 @@ def __init__(self, filename): self.stream = self.output.add_stream("h264", str(FPS)) self.stream.height = RESOLUTION[1] self.stream.width = RESOLUTION[0] - self.stream.bit_rate = 8500e3 + self.stream.bit_rate = 8500e1 def start(self): self._start = time.time() From 0adb7734b4df5773762e7ee12f61c48a3735e464 Mon Sep 17 00:00:00 2001 From: apirrone Date: Sat, 9 Sep 2023 21:06:40 +0200 Subject: [PATCH 2/7] just adding a black bar around search bar for better visibility on white frames --- memento/timeline/frame_getter.py | 9 ++++++++- memento/timeline/search_bar.py | 23 +++++++++++++++-------- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/memento/timeline/frame_getter.py b/memento/timeline/frame_getter.py index 366b1db..0e83acd 100644 --- a/memento/timeline/frame_getter.py +++ b/memento/timeline/frame_getter.py @@ -130,10 +130,17 @@ 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_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(): diff --git a/memento/timeline/search_bar.py b/memento/timeline/search_bar.py index 2b2b56c..df78cbb 100644 --- a/memento/timeline/search_bar.py +++ b/memento/timeline/search_bar.py @@ -1,6 +1,7 @@ import pygame import pygame_textinput from memento.db import Db +import numpy as np class SearchBar: @@ -8,11 +9,11 @@ def __init__(self, frame_getter): self.frame_getter = frame_getter self.db = Db() - ws = self.frame_getter.window_size - self.x = ws[0] // 10 - self.y = ws[1] // 10 - self.w = ws[0] - ws[0] // 5 - self.h = ws[1] // 20 + self.ws = self.frame_getter.window_size + self.x = self.ws[0] // 10 + self.y = self.ws[1] // 10 + self.w = self.ws[0] - self.ws[0] // 5 + self.h = self.ws[1] // 20 self.active = False self.found = False self.input_changed = False @@ -35,8 +36,15 @@ def draw_bar(self, screen): (self.x, self.y, self.w, self.h), border_radius=self.h // 4, ) + # draw a black border + pygame.draw.rect( + screen, + (0, 0, 0), + (self.x, self.y, self.w, self.h), + width=5, + border_radius=self.h // 4, + ) - def draw_text(self, screen): font = pygame.font.SysFont("Arial", self.h // 2) text = font.render(self.textinput.value, True, (0, 0, 0)) screen.blit(text, (self.x + 10, self.y + 10)) @@ -62,7 +70,6 @@ def draw(self, screen): if not self.active: return self.draw_bar(screen) - self.draw_text(screen) def start_query(self): print("start query") @@ -71,7 +78,7 @@ def start_query(self): results = self.db.search(query_input) - self.frame_getter.set_annotation(results) + self.frame_getter.set_annotations(results) if len(results) > 0: self.found = True self.frame_getter.nb_results = len(results) From f10bb6a45bb675f39e27daefb4983340e960284b Mon Sep 17 00:00:00 2001 From: apirrone Date: Sun, 10 Sep 2023 21:15:07 +0200 Subject: [PATCH 3/7] making an Apps class to avoid code duplication in the future --- memento/timeline/apps.py | 55 ++++++++++++++++++++++++++++++++ memento/timeline/search_bar.py | 3 ++ memento/timeline/time_bar.py | 57 +++++++--------------------------- 3 files changed, 69 insertions(+), 46 deletions(-) create mode 100644 memento/timeline/apps.py diff --git a/memento/timeline/apps.py b/memento/timeline/apps.py new file mode 100644 index 0000000..c35f398 --- /dev/null +++ b/memento/timeline/apps.py @@ -0,0 +1,55 @@ +import numpy as np +from memento.timeline.icon_getter import IconGetter + +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) + + for i in range(self.nb_frames): + 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"] diff --git a/memento/timeline/search_bar.py b/memento/timeline/search_bar.py index df78cbb..d13061b 100644 --- a/memento/timeline/search_bar.py +++ b/memento/timeline/search_bar.py @@ -66,6 +66,9 @@ def draw_bar(self, screen): ), ) + def draw_app_filter(self, screen): + pass + def draw(self, screen): if not self.active: return diff --git a/memento/timeline/time_bar.py b/memento/timeline/time_bar.py index 00efb84..8c5899f 100644 --- a/memento/timeline/time_bar.py +++ b/memento/timeline/time_bar.py @@ -1,29 +1,9 @@ -from memento.timeline.icon_getter import IconGetter import memento.utils as utils import cv2 import numpy as np import pygame import datetime - -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 ? +from memento.timeline.apps import Apps class TimeBar: @@ -52,13 +32,9 @@ def __init__(self, frame_getter): self.preview_frame_i = 0 self.preview_surf = None - self.ig = IconGetter(size=self.h) - - self.apps = {} + self.apps = Apps(self.frame_getter) self.today = datetime.datetime.now().strftime("%Y-%m-%d") - self.build() - def zoom(self, dir): self.tws = min( min(self.nb_frames, utils.MAX_TWS), @@ -87,19 +63,6 @@ def get_friendly_date(self, date): else: return date - def build(self): - for i in range(self.nb_frames): - 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 - # TODO adjust time bar so that the cursor is visible when jumping that way def set_current_frame_i(self, frame_i): # nb_moves = frame_i - self.current_frame_i @@ -165,7 +128,7 @@ def draw_bar(self, screen, mouse_pos): pygame.draw.rect( screen, - self.apps[app]["color"], + self.apps.get_color(app), (seg_x, self.y, seg_w, self.h), border_radius=self.h // 4, ) @@ -178,23 +141,25 @@ def draw_bar(self, screen, mouse_pos): seg_x = self.x + (start / self.tws) * self.w seg_w = (end - start) / self.tws * self.w - if self.apps[app]["icon_small"] is not None: + if self.apps.get_icon(app) is not None: if start <= self.current_frame_i <= end or utils.in_rect( (seg_x, self.y, seg_w, self.h), mouse_pos ): screen.blit( - self.apps[app]["icon_big"], + self.apps.get_icon(app, small=False), ( - self.x + (middle / self.tws) * self.w - self.ig.size, - self.y - self.ig.size // 2, + self.x + (middle / self.tws) * self.w - self.apps.ig.size, + self.y - self.apps.ig.size // 2, ), ) else: screen.blit( - self.apps[app]["icon_small"], + self.apps.get_icon(app, small=True), ( - self.x + (middle / self.tws) * self.w - self.ig.size // 2, + self.x + + (middle / self.tws) * self.w + - self.apps.ig.size // 2, self.y, ), ) From b5fc59466dd1926e173f7f8d4d2f127fc282a526 Mon Sep 17 00:00:00 2001 From: apirrone Date: Thu, 14 Sep 2023 15:02:42 +0200 Subject: [PATCH 4/7] search list --- memento/db.py | 2 +- memento/timeline/apps.py | 2 +- memento/timeline/frame_getter.py | 14 +++- memento/timeline/search_bar.py | 132 +++++++++++++++++++++++++++++-- memento/timeline/timeline.py | 32 ++++++-- 5 files changed, 165 insertions(+), 17 deletions(-) diff --git a/memento/db.py b/memento/db.py index c4173f5..3f1c508 100644 --- a/memento/db.py +++ b/memento/db.py @@ -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: diff --git a/memento/timeline/apps.py b/memento/timeline/apps.py index c35f398..1a365f4 100644 --- a/memento/timeline/apps.py +++ b/memento/timeline/apps.py @@ -37,7 +37,7 @@ def __init__(self, frame_getter): 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): + 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)) diff --git a/memento/timeline/frame_getter.py b/memento/timeline/frame_getter.py index 0e83acd..81fb61a 100644 --- a/memento/timeline/frame_getter.py +++ b/memento/timeline/frame_getter.py @@ -32,9 +32,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 +47,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 @@ -134,6 +141,9 @@ 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]: diff --git a/memento/timeline/search_bar.py b/memento/timeline/search_bar.py index d13061b..4c3c979 100644 --- a/memento/timeline/search_bar.py +++ b/memento/timeline/search_bar.py @@ -2,6 +2,9 @@ import pygame_textinput from memento.db import Db import numpy as np +import memento.timeline.text_utils as text_utils +from thefuzz import fuzz +import memento.utils as utils class SearchBar: @@ -10,14 +13,27 @@ def __init__(self, frame_getter): self.db = Db() self.ws = self.frame_getter.window_size - self.x = self.ws[0] // 10 - self.y = self.ws[1] // 10 - self.w = self.ws[0] - self.ws[0] // 5 + # self.x = self.ws[0] // 10 + # self.y = self.ws[1] // 10 + # self.w = self.ws[0] - self.ws[0] // 5 + # self.h = self.ws[1] // 20 + + self.x = 0 + self.y = 0 + self.w = self.ws[0] // 4 self.h = self.ws[1] // 20 self.active = False self.found = False self.input_changed = False + self.list_entry_h = self.ws[1] // 20 + self.list_entry_border = 5 + self.selected_entry_frame_i = None + self.font_size = 20 + self.font = pygame.font.SysFont("Arial", self.font_size) + + self.y_offset = self.h + self.textinput = pygame_textinput.TextInputManager() def activate(self): @@ -42,7 +58,7 @@ def draw_bar(self, screen): (0, 0, 0), (self.x, self.y, self.w, self.h), width=5, - border_radius=self.h // 4, + # border_radius=self.h // 4, ) font = pygame.font.SysFont("Arial", self.h // 2) @@ -66,17 +82,115 @@ def draw_bar(self, screen): ), ) + # TODO too long to compute as is. + # Should be better with segments + def remove_similar_annotations(self, annotations): + filtered_annotations = {} + annotations_frames = [] + annotations_texts = [] + for key, annotations_list in annotations.items(): + for annotation in annotations_list: + annot1 = annotation["text"] + if len(annotations_texts) == 0: + annotations_frames.append(key) + annotations_texts.append(annotation["text"]) + continue + for annot2 in annotations_texts: + if fuzz.ratio(annot1, annot2) < 30: + annotations_frames.append(key) + annotations_texts.append(annotation["text"]) + + print(annotations_frames) + + def draw_results_list(self, screen): + annotations = self.frame_getter.get_annotations() + # self.remove_similar_annotations(annotations) + + pygame.draw.rect( + screen, + (255, 255, 255), + (self.x, self.y, self.w, self.ws[1]), + border_radius=self.h // 4, + ) + # draw a black border + pygame.draw.rect( + screen, + (0, 0, 0), + (self.x, self.y, self.w, self.ws[1]), + width=5, + # border_radius=self.h // 4, + ) + index = 0 + for key, annotations_list in annotations.items(): + for annotation in annotations_list: + text = annotation["text"] + text = text.replace("\n", " ") + rect = ( + self.x, + self.y + self.y_offset + index * self.list_entry_h, + self.w, + self.list_entry_h, + ) + # y = self.y + self.y_offset + index * self.list_entry_h, + pygame.draw.rect( + screen, + (0, 0, 0), + rect, + width=2, + ) + + text_utils.render_text( + screen, + text[: self.w // 10], # TODO compute this properly + self.font, + self.x + self.list_entry_border, + self.y + + self.y_offset + + index * self.list_entry_h + + self.list_entry_border, + 1000000, + (0, 0, 0), + ) + if int(key) == int(self.frame_getter.current_displayed_frame_i): + pygame.draw.rect( + screen, + (255, 0, 0), + rect, + width=2, + ) + + if utils.in_rect(rect, pygame.mouse.get_pos()): + pygame.draw.rect( + screen, + (0, 255, 0), + rect, + width=5, + ) + self.selected_entry_frame_i = int(key) + + index += 1 + # break + + def hover(self, mouse_pos): + return utils.in_rect((self.x, self.y, self.w, self.ws[1]), mouse_pos) + + def scroll(self, dir): + if len(self.frame_getter.annotations) > 0: + self.y_offset = min(self.h, self.y_offset + dir * 10) + def draw_app_filter(self, screen): pass def draw(self, screen): if not self.active: return + self.draw_results_list(screen) self.draw_bar(screen) def start_query(self): print("start query") self.frame_getter.clear_annotations() + self.y_offset = self.h query_input = self.textinput.value results = self.db.search(query_input) @@ -91,7 +205,8 @@ def start_query(self): def events(self, events): self.textinput.update(events) - ret = False + found = False + frame_i = None for event in events: if event.type == pygame.KEYDOWN: if event.mod & pygame.KMOD_CTRL and event.key == pygame.K_f: @@ -106,8 +221,11 @@ def events(self, events): if self.input_changed or not self.found: self.start_query() if self.found: - ret = True + found = True self.input_changed = False elif self.active: self.input_changed = True - return ret + if event.type == pygame.MOUSEBUTTONDOWN: + if event.button == 1: + frame_i = self.selected_entry_frame_i + return found, frame_i diff --git a/memento/timeline/timeline.py b/memento/timeline/timeline.py index 24c0d45..fb4c68a 100644 --- a/memento/timeline/timeline.py +++ b/memento/timeline/timeline.py @@ -46,9 +46,21 @@ def update(self): self.chat = Chat(self.frame_getter) def draw_current_frame(self): - frame = self.frame_getter.get_frame(self.time_bar.current_frame_i) + x = 0 + y = 0 + resize = None + if self.search_bar.active: + w = self.window_size[0] - self.search_bar.w + ratio = self.window_size[1] / self.window_size[0] + h = int(w * ratio) + resize = (w, h) + x = self.search_bar.w + y = (self.window_size[1] - h) // 2 + frame = self.frame_getter.get_frame( + self.time_bar.current_frame_i, resize=resize + ) surf = pygame.surfarray.make_surface(frame).convert() - self.screen.blit(surf, (0, 0)) + self.screen.blit(surf, (x, y)) # TODO This is a mess def handle_inputs(self): @@ -56,7 +68,7 @@ def handle_inputs(self): ret_frame = None mouse_wheel = 0 events = pygame.event.get() - found = self.search_bar.events(events) + found, search_bar_frame_i = self.search_bar.events(events) ret_frame = self.chat.events(events) for event in events: if event.type == pygame.MOUSEWHEEL: @@ -64,10 +76,14 @@ def handle_inputs(self): if not self.ctrl_pressed: if self.chat.active and self.chat.hover(pygame.mouse.get_pos()): self.chat.scroll(event.y) + elif self.search_bar.active and self.search_bar.hover( + pygame.mouse.get_pos() + ): + self.search_bar.scroll(event.y) else: self.time_bar.move_cursor((mouse_wheel)) self.region_selector.reset() - self.frame_getter.clear_annotations() + # self.frame_getter.clear_annotations() if event.type == pygame.MOUSEBUTTONDOWN: if event.button == 1: if self.time_bar.hover(event.pos): @@ -75,6 +91,8 @@ def handle_inputs(self): self.time_bar.get_frame_i(event.pos) ) self.region_selector.reset() + elif self.search_bar.active and self.search_bar.hover(event.pos): + pass elif not self.chat.hover(pygame.mouse.get_pos()): self.search_bar.deactivate() self.region_selector.start(event.pos) @@ -150,7 +168,9 @@ def handle_inputs(self): 2, ) - if found: + if search_bar_frame_i is not None: + self.time_bar.set_current_frame_i(search_bar_frame_i) + elif found: self.time_bar.set_current_frame_i( self.frame_getter.get_next_annotated_frame_i() ) @@ -195,7 +215,7 @@ def draw_and_compute_fps(self): def run(self): while True: - self.screen.fill((255, 255, 255)) + self.screen.fill((0, 0, 0)) self.draw_current_frame() self.time_bar.draw(self.screen, pygame.mouse.get_pos()) self.search_bar.draw(self.screen) From b07ca0dbfad052963e252aab031367a9469087b8 Mon Sep 17 00:00:00 2001 From: apirrone Date: Tue, 19 Sep 2023 09:57:14 +0200 Subject: [PATCH 5/7] search list does not display out of screen items --- memento/timeline/search_bar.py | 12 +++++++++--- memento/utils.py | 6 ++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/memento/timeline/search_bar.py b/memento/timeline/search_bar.py index 4c3c979..389d76c 100644 --- a/memento/timeline/search_bar.py +++ b/memento/timeline/search_bar.py @@ -82,8 +82,8 @@ def draw_bar(self, screen): ), ) - # TODO too long to compute as is. - # Should be better with segments + # TODO too long to compute as is. + # Should be better with segments def remove_similar_annotations(self, annotations): filtered_annotations = {} annotations_frames = [] @@ -131,6 +131,12 @@ def draw_results_list(self, screen): self.w, self.list_entry_h, ) + if not utils.rect_in_rect( + rect, (0, 0, self.ws[0], self.ws[1] + rect[3]) + ): + index += 1 + continue + # y = self.y + self.y_offset + index * self.list_entry_h, pygame.draw.rect( screen, @@ -173,7 +179,7 @@ def draw_results_list(self, screen): def hover(self, mouse_pos): return utils.in_rect((self.x, self.y, self.w, self.ws[1]), mouse_pos) - + def scroll(self, dir): if len(self.frame_getter.annotations) > 0: self.y_offset = min(self.h, self.y_offset + dir * 10) diff --git a/memento/utils.py b/memento/utils.py index 825914f..7eac9fb 100644 --- a/memento/utils.py +++ b/memento/utils.py @@ -98,6 +98,12 @@ def in_rect(rect, pos): return x <= pos[0] <= x + w and y <= pos[1] <= y + h +def rect_in_rect(rect1, rect2): + x1, y1, w1, h1 = rect1 + x2, y2, w2, h2 = rect2 + return x1 >= x2 and y1 >= y2 and x1 + w1 <= x2 + w2 and y1 + h1 <= y2 + h2 + + def draw_results(res, frame): for entry in res: x = int(entry["x"]) From 9ca0c77e60b7b913ac1b783fec2caf6dea21035e Mon Sep 17 00:00:00 2001 From: apirrone Date: Tue, 19 Sep 2023 14:34:40 +0200 Subject: [PATCH 6/7] new search UI + faster timeline loading time + housekeeping --- memento/background.py | 18 +++++++------- memento/caching.py | 14 +++++------ memento/db.py | 4 ++-- memento/timeline/apps.py | 7 +++++- memento/timeline/chat.py | 16 +++++++++---- memento/timeline/frame_getter.py | 7 +++--- memento/timeline/icon_getter.py | 40 +++++++++++++++++++++++--------- memento/timeline/time_bar.py | 1 - memento/utils.py | 1 + 9 files changed, 67 insertions(+), 41 deletions(-) diff --git a/memento/background.py b/memento/background.py index 0aeabbf..20a3bac 100644 --- a/memento/background.py +++ b/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") ) diff --git a/memento/caching.py b/memento/caching.py index f96dddf..6322240 100644 --- a/memento/caching.py +++ b/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) diff --git a/memento/db.py b/memento/db.py index 3f1c508..e3bf24a 100644 --- a/memento/db.py +++ b/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 diff --git a/memento/timeline/apps.py b/memento/timeline/apps.py index 1a365f4..97d41a8 100644 --- a/memento/timeline/apps.py +++ b/memento/timeline/apps.py @@ -1,5 +1,6 @@ import numpy as np from memento.timeline.icon_getter import IconGetter +import time COLOR_PALETTE = [ [244, 67, 54], @@ -33,7 +34,11 @@ def __init__(self, frame_getter): self.nb_frames = self.frame_getter.nb_frames self.ig = IconGetter(size=self.h) - for i in range(self.nb_frames): + # 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] = {} diff --git a/memento/timeline/chat.py b/memento/timeline/chat.py index 9492667..5419d7a 100644 --- a/memento/timeline/chat.py +++ b/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", ) @@ -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, @@ -130,6 +134,8 @@ def process_chat_query(self): 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 +350,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) diff --git a/memento/timeline/frame_getter.py b/memento/timeline/frame_getter.py index 81fb61a..f94ea31 100644 --- a/memento/timeline/frame_getter.py +++ b/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 ) diff --git a/memento/timeline/icon_getter.py b/memento/timeline/icon_getter.py index 4255f72..ca3adde 100644 --- a/memento/timeline/icon_getter.py +++ b/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 diff --git a/memento/timeline/time_bar.py b/memento/timeline/time_bar.py index 8c5899f..e41d714 100644 --- a/memento/timeline/time_bar.py +++ b/memento/timeline/time_bar.py @@ -31,7 +31,6 @@ def __init__(self, frame_getter): self.current_frame_i = self.tw_end - 1 self.preview_frame_i = 0 self.preview_surf = None - self.apps = Apps(self.frame_getter) self.today = datetime.datetime.now().strftime("%Y-%m-%d") diff --git a/memento/utils.py b/memento/utils.py index 7eac9fb..609eb3e 100644 --- a/memento/utils.py +++ b/memento/utils.py @@ -17,6 +17,7 @@ RESOLUTION = (1920, 1080) MAX_TWS = 10000 * FPS * SECONDS_PER_REC FRAME_CACHE_SIZE = int((MAX_TWS / FPS / SECONDS_PER_REC)) +CACHE_PATH = os.path.join(os.environ["HOME"], ".cache", "memento") def get_active_window(): From 3b964b0b3eb318e1c53da1dae65a09c30d8b0a21 Mon Sep 17 00:00:00 2001 From: apirrone Date: Wed, 20 Sep 2023 13:35:45 +0200 Subject: [PATCH 7/7] Adding try catch when trying to decode the llm's json output. Changing the prompt a little bit too --- memento/timeline/chat.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/memento/timeline/chat.py b/memento/timeline/chat.py index 5419d7a..f969333 100644 --- a/memento/timeline/chat.py +++ b/memento/timeline/chat.py @@ -65,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. @@ -126,7 +126,13 @@ 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)