feat: Refactor The Finals

crispy-api/api/__init__.py

@@ -12,27 +12,27 @@
 from montydb import MontyClient, set_storage
 from pydantic.json import ENCODERS_BY_TYPE
 
-from api.config import ASSETS, DATABASE_PATH, DEBUG, FRAMERATE, GAME, MUSICS, VIDEOS
+from api.config import (
+    ASSETS,
+    DATABASE_PATH,
+    DEBUG,
+    FRAMERATE,
+    GAME,
+    MUSICS,
+    USE_NETWORK,
+    VIDEOS,
+)
 from api.tools.AI.network import NeuralNetwork
-from api.tools.enums import SupportedGames
 from api.tools.filters import apply_filters  # noqa
 from api.tools.setup import handle_highlights, handle_musics
 
 ENCODERS_BY_TYPE[ObjectId] = str
 
-neural_network = None
-
-if GAME != SupportedGames.THEFINALS:
-    neural_network = NeuralNetwork(GAME)
+NEURAL_NETWORK = None
 
-    if GAME == SupportedGames.OVERWATCH:
-        neural_network.load(os.path.join(ASSETS, "overwatch.npy"))
-    elif GAME == SupportedGames.VALORANT:
-        neural_network.load(os.path.join(ASSETS, "valorant.npy"))
-    elif GAME == SupportedGames.CSGO2:
-        neural_network.load(os.path.join(ASSETS, "csgo2.npy"))
-    else:
-        raise ValueError(f"game {GAME} not supported")
+if USE_NETWORK:
+    NEURAL_NETWORK = NeuralNetwork(GAME)
+    NEURAL_NETWORK.load(os.path.join(ASSETS, GAME + ".npy"))
 
 
 logging.getLogger("PIL").setLevel(logging.ERROR)

crispy-api/api/config.py

@@ -64,3 +64,4 @@
         raise ValueError(f"game {GAME} not supported")
 
 READER = easyocr.Reader(["fr", "en"], gpu=True, verbose=False)
+USE_NETWORK = GAME not in [SupportedGames.THEFINALS]

crispy-api/api/models/highlight.py

@@ -116,7 +116,6 @@ async def extract_images(
         images_path = os.path.join(self.directory, save_path)
 
         if not os.path.exists(images_path):
-            print("creating images path at", images_path)
             os.mkdir(images_path)
             (
                 ffmpeg.input(self.path)

crispy-api/api/routes/highlight.py

@@ -5,7 +5,7 @@
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
 
-from api import app, neural_network
+from api import NEURAL_NETWORK, app
 from api.config import CONFIDENCE, FRAMERATE, FRAMES_AFTER, FRAMES_BEFORE, OFFSET
 from api.models.highlight import Highlight
 from api.models.segment import Segment
@@ -84,7 +84,7 @@ async def post_highlights_segments_generate() -> None:
             extract_segments,
             kwargs={
                 "highlight": highlight,
-                "neural_network": neural_network,
+                "neural_network": NEURAL_NETWORK,
                 "confidence": CONFIDENCE,
                 "framerate": FRAMERATE,
                 "offset": OFFSET,

crispy-api/api/tools/setup.py

@@ -24,6 +24,41 @@ def __sanitize_path(path: str) -> str:
     return path
 
 
+def handle_the_finals(
+    new_highlights: List[Highlight],
+    framerate: int = 4,
+) -> None:
+    for highlight in new_highlights:
+        path = os.path.join(highlight.directory, "usernames")
+        images = os.listdir(path)
+        usernames: List[str] = ["", ""]
+        usernames_histogram: Counter = Counter()
+
+        for i in range(0, len(images), framerate):
+            image = images[i]
+            image_path = os.path.join(path, image)
+            result = READER.readtext(image_path)
+            for text in result:
+                if text[1].isnumeric():
+                    continue
+                usernames_histogram[text[1]] += 1
+            most_common_usernames = usernames_histogram.most_common(2)
+            if most_common_usernames[0][1] >= 10 and most_common_usernames[1][1] >= 10:
+                usernames = [most_common_usernames[0][0], most_common_usernames[1][0]]
+                break
+        highlight.update({"usernames": usernames})
+        highlight.save()
+
+
+def handle_specific_game(
+    new_highlights: List[Highlight],
+    game: SupportedGames,
+    framerate: int = 4,
+) -> None:
+    if game == SupportedGames.THEFINALS:
+        handle_the_finals(new_highlights, framerate)
+
+
 async def handle_highlights(
     path: str,
     game: SupportedGames,
@@ -105,30 +140,7 @@ async def handle_highlights(
 
     Highlight.update_many({}, {"$set": {"job_id": None}})
 
-    if game == SupportedGames.THEFINALS:
-        path = os.path.join(highlight.directory, "usernames")
-        for highlight in new_highlights:
-            images = os.listdir(path)
-            usernames = [""] * 2
-            usernames_histogram: Counter = Counter()
-
-            for i in range(0, len(images), framerate):
-                image = images[i]
-                image_path = os.path.join(path, image)
-                result = READER.readtext(image_path)
-                for text in result:
-                    if text[1].isnumeric():
-                        continue
-                    usernames_histogram[text[1]] += 1
-                two_best = usernames_histogram.most_common(2)
-                if two_best[0][1] >= 10 and two_best[1][1] >= 10:
-                    usernames = [
-                        usernames_histogram.most_common(2)[0][0],
-                        usernames_histogram.most_common(2)[1][0],
-                    ]
-                    break
-            highlight.update({"usernames": usernames})
-            highlight.save()
+    handle_specific_game(new_highlights, game, framerate)
 
     return new_highlights
 

crispy-api/api/tools/video.py

@@ -56,7 +56,7 @@ def _create_query_array(
     return queries
 
 
-def _get_the_finals_query_array(highlight: Highlight) -> List[int]:
+def _create_the_finals_query_array(highlight: Highlight) -> List[int]:
     usernames = highlight.usernames
     images = os.listdir(highlight.images_path)
     images.sort()
@@ -80,7 +80,7 @@ def _get_query_array(
     if neural_network:
         return _create_query_array(neural_network, highlight, confidence)
     if GAME == SupportedGames.THEFINALS:
-        return _get_the_finals_query_array(highlight)
+        return _create_the_finals_query_array(highlight)
     raise ValueError(f"No neural network for game {GAME} and no custom query array")
 
 
@@ -158,7 +158,6 @@ async def extract_segments(
     Extract segments from a highlight
 
     :param highlight: highlight to extract segments from
-    :param neural_network: neural network to query
     :param confidence: confidence to query
     :param offset: offset to post process
     :param framerate: framerate of the video