mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-14 14:01:27 +00:00
Sleeps
62 lines
2.7 KiB
Python
62 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from unittest.mock import Mock
|
|
|
|
from scene_search.caption_text import remove_image_preface
|
|
from scene_search.captioner import OllamaCaptioner
|
|
|
|
|
|
class CaptionTextTests(unittest.TestCase):
|
|
def test_removes_model_image_set_introductions(self) -> None:
|
|
examples = {
|
|
"The video features a series of still images showing a man wearing a red coat.":
|
|
"A man wearing a red coat.",
|
|
"A series of images depicts a woman standing beside a blue car.":
|
|
"A woman standing beside a blue car.",
|
|
"These frames show two people walking along a beach.":
|
|
"Two people walking along a beach.",
|
|
"This clip contains a collection of photographs depicting an old bridge at night.":
|
|
"An old bridge at night.",
|
|
"The images depict a small dog lying on the sofa.":
|
|
"A small dog lying on the sofa.",
|
|
}
|
|
for original, expected in examples.items():
|
|
with self.subTest(original=original):
|
|
self.assertEqual(remove_image_preface(original), expected)
|
|
|
|
def test_only_anchored_presentation_boilerplate_is_removed(self) -> None:
|
|
examples = [
|
|
"A photographer is sorting a series of photos on a desk.",
|
|
"The video features a series of still images showing a",
|
|
"A row of framed pictures hangs behind a woman.",
|
|
"A person displays a sequence of images on a monitor.",
|
|
"A woman walks through a building and waves to the camera.",
|
|
]
|
|
for original in examples:
|
|
with self.subTest(original=original):
|
|
self.assertEqual(remove_image_preface(original), original)
|
|
|
|
def test_live_caption_generation_normalises_response_and_changes_prompt(self) -> None:
|
|
captioner = OllamaCaptioner("http://ollama.test", "vision-test")
|
|
response = Mock()
|
|
response.ok = True
|
|
response.json.return_value = {
|
|
"message": {
|
|
"content": (
|
|
"The video features a series of still images showing a "
|
|
"person walking across a street."
|
|
)
|
|
}
|
|
}
|
|
captioner.session.post = Mock(return_value=response)
|
|
result = captioner._caption_b64("an-image", chronological_sequence=True)
|
|
self.assertEqual(result, "A person walking across a street.")
|
|
prompt = captioner.session.post.call_args.kwargs["json"]["messages"][0]["content"]
|
|
self.assertIn("Never introduce your answer", prompt)
|
|
self.assertIn("Do not invent motion", prompt)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|