From 42e3473f28d8c400c49c6c159d7d3d0367450ca0 Mon Sep 17 00:00:00 2001 From: Flowtter Date: Fri, 5 Jan 2024 11:48:57 +0100 Subject: [PATCH] feat: Add support for 4/3 --- README.md | 5 ++ crispy-api/api/__init__.py | 9 ++-- crispy-api/api/config.py | 1 + crispy-api/api/models/highlight.py | 83 +++++++++++++++++++++++++----- crispy-api/api/tools/ffmpeg.py | 2 +- crispy-api/api/tools/setup.py | 39 +++++++------- crispy-api/settings.json | 1 + crispy-api/tests/assets | 2 +- crispy-api/tests/constants.py | 2 + crispy-api/tests/tools/setup.py | 30 ++++++++++- etc/default-settings.json | 1 + 11 files changed, 134 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 55d5284..10cde67 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ config example: "second-after": 3, "second-between-kills": 5 }, + "stretch": false, "game": "valorant" } ``` @@ -60,6 +61,7 @@ The following settings are adjustable: - second-after: Seconds of gameplay included after the highlight. - second-between-kills: Transition time between highlights. If the time between two highlights is less than this value, the both highlights will be merged. - game: Chosen game (either "valorant", "overwatch" or "csgo2") +- stretch: This is an option in case you're playing on a 4:3 resolution but your clips are recorded in 16:9. ### Recommended settings @@ -79,6 +81,7 @@ Here are some settings that I found to work well for me: "second-after": 0.5, "second-between-kills": 3 }, + "stretch": false, "game": "valorant" } ``` @@ -96,6 +99,7 @@ Here are some settings that I found to work well for me: "second-after": 3, "second-between-kills": 5 }, + "stretch": false, "game": "overwatch" } ``` @@ -113,6 +117,7 @@ Here are some settings that I found to work well for me: "second-after": 1, "second-between-kills": 3 }, + "stretch": false, "game": "csgo2" } ``` diff --git a/crispy-api/api/__init__.py b/crispy-api/api/__init__.py index 33a2f49..af3535c 100644 --- a/crispy-api/api/__init__.py +++ b/crispy-api/api/__init__.py @@ -1,4 +1,5 @@ import logging +import os import subprocess from typing import Optional @@ -11,7 +12,7 @@ from montydb import MontyClient, set_storage from pydantic.json import ENCODERS_BY_TYPE -from api.config import DATABASE_PATH, DEBUG, FRAMERATE, GAME, MUSICS, VIDEOS +from api.config import ASSETS, DATABASE_PATH, DEBUG, FRAMERATE, GAME, MUSICS, VIDEOS from api.tools.AI.network import NeuralNetwork from api.tools.enums import SupportedGames from api.tools.filters import apply_filters # noqa @@ -22,11 +23,11 @@ neural_network = NeuralNetwork(GAME) if GAME == SupportedGames.OVERWATCH: - neural_network.load("./assets/overwatch.npy") + neural_network.load(os.path.join(ASSETS, "overwatch.npy")) elif GAME == SupportedGames.VALORANT: - neural_network.load("./assets/valorant.npy") + neural_network.load(os.path.join(ASSETS, "valorant.npy")) elif GAME == SupportedGames.CSGO2: - neural_network.load("./assets/csgo2.npy") + neural_network.load(os.path.join(ASSETS, "csgo2.npy")) else: raise ValueError(f"game {GAME} not supported") diff --git a/crispy-api/api/config.py b/crispy-api/api/config.py index 8efca99..ea8eaf4 100644 --- a/crispy-api/api/config.py +++ b/crispy-api/api/config.py @@ -55,6 +55,7 @@ CONFIDENCE = __neural_network.get("confidence", 0.6) + STRETCH = __settings.get("stretch", False) GAME = __settings.get("game") if GAME is None: raise KeyError("game not found in settings.json") diff --git a/crispy-api/api/models/highlight.py b/crispy-api/api/models/highlight.py index c9b0a34..969c60c 100644 --- a/crispy-api/api/models/highlight.py +++ b/crispy-api/api/models/highlight.py @@ -19,6 +19,42 @@ csgo2_mask = Image.open(CSGO2_MASK_PATH) +class Box: + def __init__( + self, + offset_x: int, + y: int, + width: int, + height: int, + shift_x: int, + stretch: bool, + ) -> None: + """ + :param offset_x: Offset in pixels from the center of the video to the left + :param y: Offset in pixels from the top of the video + :param width: Width of the box in pixels + :param height: Height of the box in pixels + :param shift_x: Shift the box by a certain amount of pixels to the right + + example: + If you want to create a box at 50 px from the center on x, but shifted by 20px to the right + you would do: + Box(50, 0, 100, 100, 20) + """ + half = 720 if stretch else 960 + + self.x = half - offset_x + shift_x + self.y = y + self.width = width + self.height = height + + def __iter__(self) -> Any: + yield self.x + yield self.y + yield self.width + yield self.height + + class Highlight(Thingy): segments_path: Optional[str] local_filters: Optional[Dict[str, Any]] @@ -55,7 +91,7 @@ async def extract_thumbnails(self) -> bool: async def extract_images( self, post_process: Callable, - coordinates: Tuple, + coordinates: Box, framerate: int = 4, ) -> bool: """ @@ -93,7 +129,9 @@ async def extract_images( return True - async def extract_overwatch_images(self, framerate: int = 4) -> bool: + async def extract_overwatch_images( + self, framerate: int = 4, stretch: bool = False + ) -> bool: def post_process(image: Image) -> Image: r, g, b = image.split() for x in range(image.width): @@ -121,10 +159,12 @@ def post_process(image: Image) -> Image: return final return await self.extract_images( - post_process, (910, 490, 100, 100), framerate=framerate + post_process, Box(50, 490, 100, 100, 0, stretch), framerate=framerate ) - async def extract_valorant_images(self, framerate: int = 4) -> bool: + async def extract_valorant_images( + self, framerate: int = 4, stretch: bool = False + ) -> bool: def _apply_filter_and_do_operations( image: Image, image_filter: ImageFilter ) -> Image: @@ -159,10 +199,12 @@ def post_process(image: Image) -> Image: return final return await self.extract_images( - post_process, (899, 801, 122, 62), framerate=framerate + post_process, Box(61, 801, 122, 62, 0, stretch), framerate=framerate ) - async def extract_csgo2_images(self, framerate: int = 4) -> bool: + async def extract_csgo2_images( + self, framerate: int = 4, stretch: bool = False + ) -> bool: def post_process(image: Image) -> Image: image = ImageOps.grayscale( image.filter(ImageFilter.FIND_EDGES).filter( @@ -175,18 +217,18 @@ def post_process(image: Image) -> Image: return final return await self.extract_images( - post_process, (930, 925, 100, 100), framerate=framerate + post_process, Box(50, 925, 100, 100, 20, stretch), framerate=framerate ) async def extract_images_from_game( - self, game: SupportedGames, framerate: int = 4 + self, game: SupportedGames, framerate: int = 4, stretch: bool = False ) -> bool: if game == SupportedGames.OVERWATCH: - return await self.extract_overwatch_images(framerate) + return await self.extract_overwatch_images(framerate, stretch) elif game == SupportedGames.VALORANT: - return await self.extract_valorant_images(framerate) + return await self.extract_valorant_images(framerate, stretch) elif game == SupportedGames.CSGO2: - return await self.extract_csgo2_images(framerate) + return await self.extract_csgo2_images(framerate, stretch) else: raise NotImplementedError @@ -212,7 +254,7 @@ def recompile(self) -> bool: return result async def extract_segments( - self, timestamps: List[Tuple[float, float]] + self, timestamps: List[Tuple[float, float]], stretch: bool = False ) -> List[Segment]: """ Segment a video into multiple videos @@ -269,11 +311,13 @@ async def extract_segments( for (start, end) in new_timestamps: segment_save_path = os.path.join(self.segments_path, f"{start}-{end}.mp4") + dar = 4 / 3 if stretch else 16 / 9 ( ffmpeg.input( self.path, ) .apply_filters(self.id) + .filter("setdar", dar) .output(audio, segment_save_path, ss=f"{start}", to=f"{end}") .overwrite_output() .run(quiet=True) @@ -318,7 +362,11 @@ async def concatenate_segments(self) -> bool: return True async def scale_video( - self, width: int = 1920, height: int = 1080, backup: str = BACKUP + self, + width: int = 1920, + height: int = 1080, + backup: str = BACKUP, + stretch: bool = False, ) -> None: """ Scale (up or down) a video. @@ -344,6 +392,9 @@ async def scale_video( audio = silence_if_no_audio(video.audio, backup_path) video = video.filter("scale", w=width, h=height) + video = ( + video.filter("setdar", 4 / 3) if stretch else video.filter("setdar", 16 / 9) + ) ffmpeg.output(video, audio, self.path, start_number=0).run(quiet=True) @@ -358,7 +409,7 @@ async def remove(self) -> None: Filter.delete_one({"highlight_id": self.id}) self.delete() - async def extract_snippet_in_lower_resolution(self) -> bool: + async def extract_snippet_in_lower_resolution(self, stretch: bool = False) -> bool: """Extract 5 seconds of a highlight in lower resolution""" if self.snippet_path: return False @@ -367,7 +418,11 @@ async def extract_snippet_in_lower_resolution(self) -> bool: video = ffmpeg.input(self.path, sseof="-20") audio = silence_if_no_audio(video.audio, self.path) + video = video.filter("scale", 640, -1) + video = ( + video.filter("setdar", 4 / 3) if stretch else video.filter("setdar", 16 / 9) + ) ffmpeg.output(video, audio, snippet_path, t="00:00:5").run(quiet=True) self.update({"snippet_path": snippet_path}) diff --git a/crispy-api/api/tools/ffmpeg.py b/crispy-api/api/tools/ffmpeg.py index 17f12ac..30f38d7 100644 --- a/crispy-api/api/tools/ffmpeg.py +++ b/crispy-api/api/tools/ffmpeg.py @@ -24,7 +24,7 @@ async def merge_videos( else: clips = [] for filename in videos_path: - clips.append(mpe.VideoFileClip(filename)) + clips.append(mpe.VideoFileClip(filename).resize((1920, 1080))) final_clip = mpe.concatenate_videoclips(clips) diff --git a/crispy-api/api/tools/setup.py b/crispy-api/api/tools/setup.py index b45ed94..8798055 100644 --- a/crispy-api/api/tools/setup.py +++ b/crispy-api/api/tools/setup.py @@ -1,4 +1,3 @@ -import asyncio import logging import os import shutil @@ -7,7 +6,7 @@ import ffmpeg from PIL import Image -from api.config import SESSION, SILENCE_PATH +from api.config import SESSION, SILENCE_PATH, STRETCH from api.models.filter import Filter from api.models.highlight import Highlight from api.models.music import Music @@ -29,6 +28,7 @@ async def handle_highlights( game: SupportedGames, framerate: int = 4, session: str = SESSION, + stretch: bool = STRETCH, ) -> List[Highlight]: if not os.path.exists(session): @@ -45,7 +45,6 @@ async def handle_highlights( logger.info(f"Removing highlight {highlight.path}") await highlight.remove() - job_scheduler = JobScheduler(4) new_highlights = [] for file in sorted(os.listdir(path)): file_path = os.path.join(path, file) @@ -67,20 +66,14 @@ async def handle_highlights( ).save() Filter({"highlight_id": highlight.id}).save() new_highlights.append(highlight) + await highlight.extract_thumbnails() - job_scheduler.schedule( - highlight.extract_images_from_game, - kwargs={"game": game, "framerate": framerate}, - ) - job_scheduler.schedule(highlight.extract_thumbnails) - job_scheduler.schedule(highlight.extract_snippet_in_lower_resolution) index += 1 logger.info(f"Adding {len(new_highlights)} highlights, this may take a while.") logger.warning("Wait for `Application startup complete.` to use Crispy.") - job_scheduler.run_in_thread().join() - + target_size = (1440, 1080) if stretch else (1920, 1080) for highlight in new_highlights: if not video_has_audio(highlight.path): tmp_path = os.path.join(highlight.directory, "tmp.mp4") @@ -90,18 +83,24 @@ async def handle_highlights( ffmpeg.input(highlight.path).output( video, audio, tmp_path, vcodec="copy", acodec="aac" - ).overwrite_output().run() + ).overwrite_output().run(quiet=True) shutil.move(tmp_path, highlight.path) - if Image.open(highlight.thumbnail_path_full_size).size != (1920, 1080): - await highlight.scale_video() - coroutines = [ - highlight.extract_thumbnails(), - highlight.extract_snippet_in_lower_resolution(), - highlight.extract_images_from_game(game, framerate), - ] + if Image.open(highlight.thumbnail_path_full_size).size != target_size: + await highlight.scale_video(*target_size, stretch=stretch) + await highlight.extract_thumbnails() - await asyncio.gather(*coroutines) + job_scheduler = JobScheduler(4) + for highlight in new_highlights: + job_scheduler.schedule( + highlight.extract_images_from_game, + kwargs={"game": game, "framerate": framerate, "stretch": stretch}, + ) + job_scheduler.schedule( + highlight.extract_snippet_in_lower_resolution, kwargs={"stretch": stretch} + ) + + job_scheduler.run_in_thread().join() Highlight.update_many({}, {"$set": {"job_id": None}}) diff --git a/crispy-api/settings.json b/crispy-api/settings.json index 59e000c..cefc18a 100644 --- a/crispy-api/settings.json +++ b/crispy-api/settings.json @@ -8,5 +8,6 @@ "second-after": 0.5, "second-between-kills": 1 }, + "stretch": false, "game": "valorant" } diff --git a/crispy-api/tests/assets b/crispy-api/tests/assets index 2f71a93..0c51f43 160000 --- a/crispy-api/tests/assets +++ b/crispy-api/tests/assets @@ -1 +1 @@ -Subproject commit 2f71a9352c02fbe4bbbb7f01723ec1abdfa6f569 +Subproject commit 0c51f4332690fde477f94d1ecf6e6e08dcb82903 diff --git a/crispy-api/tests/constants.py b/crispy-api/tests/constants.py index 171ac89..bdc5183 100644 --- a/crispy-api/tests/constants.py +++ b/crispy-api/tests/constants.py @@ -10,6 +10,8 @@ MAIN_VIDEO = os.path.join(VIDEOS_PATH, "main-video.mp4") MAIN_VIDEO_DOWNSCALED = os.path.join(VIDEOS_PATH, "main-video_downscaled.mp4") MAIN_VIDEO_NO_AUDIO = os.path.join(VIDEOS_PATH, "main-video-no-audio.mp4") +MAIN_VIDEO_STRETCH = os.path.join(VIDEOS_PATH, "main-video-stretch.mp4") +MAIN_VIDEO_1440 = os.path.join(VIDEOS_PATH, "main-video-1440.mp4") MAIN_VIDEO_OVERWATCH = os.path.join(VIDEOS_PATH, "main-video-overwatch.mp4") MAIN_VIDEO_CSGO2 = os.path.join(VIDEOS_PATH, "main-video-csgo2.mp4") MAIN_SEGMENT = os.path.join(VIDEOS_PATH, "main-video-segment.mp4") diff --git a/crispy-api/tests/tools/setup.py b/crispy-api/tests/tools/setup.py index 0069e26..66d992d 100644 --- a/crispy-api/tests/tools/setup.py +++ b/crispy-api/tests/tools/setup.py @@ -5,7 +5,12 @@ from api.models.music import Music from api.tools.enums import SupportedGames from api.tools.setup import handle_highlights, handle_musics -from tests.constants import MAIN_MUSIC, MAIN_VIDEO, MAIN_VIDEO_NO_AUDIO +from tests.constants import ( + MAIN_MUSIC, + MAIN_VIDEO, + MAIN_VIDEO_NO_AUDIO, + MAIN_VIDEO_STRETCH, +) async def test_handle_highlights(tmp_path): @@ -57,6 +62,29 @@ async def test_handle_highlights(tmp_path): shutil.rmtree(tmp_resources) +async def test_handle_highlights_stretch(tmp_path): + tmp_session = os.path.join(tmp_path, "session") + tmp_resources = os.path.join(tmp_path, "resources") + os.mkdir(tmp_resources) + + shutil.copy(MAIN_VIDEO_STRETCH, tmp_resources) + + assert await handle_highlights( + tmp_resources, SupportedGames.VALORANT, session=tmp_session, stretch=True + ) + + assert Highlight.count_documents() == 1 + + basename_no_ext = os.path.splitext(os.path.basename(MAIN_VIDEO_STRETCH))[0] + + shutil.copytree( + os.path.join(tmp_session, basename_no_ext, "images"), + os.path.join(tmp_path, "images"), + ) + shutil.rmtree(tmp_session) + shutil.rmtree(tmp_resources) + + async def test_handle_musics(tmp_path): tmp_resources = os.path.join(tmp_path, "resources") os.mkdir(tmp_resources) diff --git a/etc/default-settings.json b/etc/default-settings.json index 59e000c..cefc18a 100644 --- a/etc/default-settings.json +++ b/etc/default-settings.json @@ -8,5 +8,6 @@ "second-after": 0.5, "second-between-kills": 1 }, + "stretch": false, "game": "valorant" }