Spaces:
Running
on
Zero
Running
on
Zero
added text processing and tests
Browse files- src/app.py +13 -2
- tests/{test_video.py → test_media.py} +135 -1
src/app.py
CHANGED
|
@@ -63,7 +63,6 @@ def process_video(video_path: str, max_images: int) -> list[dict]:
|
|
| 63 |
result_content = []
|
| 64 |
# TODO: Change max_image to slider
|
| 65 |
frames = get_frames(video_path, max_images)
|
| 66 |
-
# Take frame and attach to result_content with timestamp
|
| 67 |
for frame in frames:
|
| 68 |
image, timestamp = frame
|
| 69 |
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
|
|
@@ -71,4 +70,16 @@ def process_video(video_path: str, max_images: int) -> list[dict]:
|
|
| 71 |
result_content.append({"type": "text", "text": f"Frame {timestamp}:"})
|
| 72 |
result_content.append({"type": "image", "url": temp_file.name})
|
| 73 |
logger.debug(f"Processed {len(frames)} frames from video {video_path} with frames {result_content}")
|
| 74 |
-
return result_content
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 63 |
result_content = []
|
| 64 |
# TODO: Change max_image to slider
|
| 65 |
frames = get_frames(video_path, max_images)
|
|
|
|
| 66 |
for frame in frames:
|
| 67 |
image, timestamp = frame
|
| 68 |
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
|
|
|
|
| 70 |
result_content.append({"type": "text", "text": f"Frame {timestamp}:"})
|
| 71 |
result_content.append({"type": "image", "url": temp_file.name})
|
| 72 |
logger.debug(f"Processed {len(frames)} frames from video {video_path} with frames {result_content}")
|
| 73 |
+
return result_content
|
| 74 |
+
|
| 75 |
+
def process_user_input(message: dict, max_images: int) -> list[dict]:
|
| 76 |
+
if not message["files"]:
|
| 77 |
+
return [{"type": "text", "text": message["text"]}]
|
| 78 |
+
|
| 79 |
+
if message["files"][0].endswith(".mp4"):
|
| 80 |
+
return [{"type": "text", "text": message["text"]}, *process_video(message["files"][0], max_images)]
|
| 81 |
+
|
| 82 |
+
return [
|
| 83 |
+
{"type": "text", "text": message["text"]},
|
| 84 |
+
*[{"type": "image", "url": path} for path in message["files"]],
|
| 85 |
+
]
|
tests/{test_video.py → test_media.py}
RENAMED
|
@@ -5,7 +5,7 @@ from PIL import Image
|
|
| 5 |
from pathlib import Path
|
| 6 |
import tempfile
|
| 7 |
|
| 8 |
-
from src.app import get_frames, process_video
|
| 9 |
|
| 10 |
# Get the project root directory
|
| 11 |
ROOT_DIR = Path(__file__).parent.parent
|
|
@@ -100,3 +100,137 @@ def test_process_video_invalid_path():
|
|
| 100 |
|
| 101 |
with pytest.raises(ValueError):
|
| 102 |
process_video("nonexistent_video.mp4", 3)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 5 |
from pathlib import Path
|
| 6 |
import tempfile
|
| 7 |
|
| 8 |
+
from src.app import get_frames, process_video, process_user_input
|
| 9 |
|
| 10 |
# Get the project root directory
|
| 11 |
ROOT_DIR = Path(__file__).parent.parent
|
|
|
|
| 100 |
|
| 101 |
with pytest.raises(ValueError):
|
| 102 |
process_video("nonexistent_video.mp4", 3)
|
| 103 |
+
|
| 104 |
+
def test_process_user_input_text_only():
|
| 105 |
+
"""Test processing user input with text only (no files)."""
|
| 106 |
+
message = {
|
| 107 |
+
"text": "This is a test message",
|
| 108 |
+
"files": []
|
| 109 |
+
}
|
| 110 |
+
|
| 111 |
+
# Add the max_images parameter
|
| 112 |
+
result = process_user_input(message, 5)
|
| 113 |
+
|
| 114 |
+
# Should return a single text item
|
| 115 |
+
assert len(result) == 1
|
| 116 |
+
assert result[0]["type"] == "text"
|
| 117 |
+
assert result[0]["text"] == "This is a test message"
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def test_process_user_input_with_video():
|
| 121 |
+
"""Test processing user input with a video file."""
|
| 122 |
+
video_path = os.path.join(ROOT_DIR, "assets", "test_video.mp4")
|
| 123 |
+
assert os.path.exists(video_path), f"Test video not found at {video_path}"
|
| 124 |
+
|
| 125 |
+
message = {
|
| 126 |
+
"text": "Video analysis",
|
| 127 |
+
"files": [video_path]
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
result = process_user_input(message, 4)
|
| 131 |
+
|
| 132 |
+
# Should have at least 3 items (text + at least one frame with text and image)
|
| 133 |
+
assert len(result) >= 3
|
| 134 |
+
|
| 135 |
+
# First item should be the message text
|
| 136 |
+
assert result[0]["type"] == "text"
|
| 137 |
+
assert result[0]["text"] == "Video analysis"
|
| 138 |
+
|
| 139 |
+
# Following items should be frame text and images
|
| 140 |
+
assert result[1]["type"] == "text"
|
| 141 |
+
assert result[1]["text"].startswith("Frame ")
|
| 142 |
+
|
| 143 |
+
assert result[2]["type"] == "image"
|
| 144 |
+
assert "url" in result[2]
|
| 145 |
+
assert os.path.exists(result[2]["url"])
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
def test_process_user_input_with_images():
|
| 149 |
+
"""Test processing user input with image files."""
|
| 150 |
+
# Create temporary image files for testing
|
| 151 |
+
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as img1, \
|
| 152 |
+
tempfile.NamedTemporaryFile(suffix=".png", delete=False) as img2:
|
| 153 |
+
|
| 154 |
+
image_paths = [img1.name, img2.name]
|
| 155 |
+
|
| 156 |
+
message = {
|
| 157 |
+
"text": "Image analysis",
|
| 158 |
+
"files": image_paths
|
| 159 |
+
}
|
| 160 |
+
|
| 161 |
+
result = process_user_input(message, 5)
|
| 162 |
+
|
| 163 |
+
# Should have 3 items (text + 2 images)
|
| 164 |
+
assert len(result) == 3
|
| 165 |
+
|
| 166 |
+
# First item should be the message text
|
| 167 |
+
assert result[0]["type"] == "text"
|
| 168 |
+
assert result[0]["text"] == "Image analysis"
|
| 169 |
+
|
| 170 |
+
# Following items should be images
|
| 171 |
+
assert result[1]["type"] == "image"
|
| 172 |
+
assert result[1]["url"] == image_paths[0]
|
| 173 |
+
|
| 174 |
+
assert result[2]["type"] == "image"
|
| 175 |
+
assert result[2]["url"] == image_paths[1]
|
| 176 |
+
|
| 177 |
+
# Clean up temp files
|
| 178 |
+
for path in image_paths:
|
| 179 |
+
if os.path.exists(path):
|
| 180 |
+
os.unlink(path)
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def test_process_user_input_empty_text():
|
| 184 |
+
"""Test processing user input with empty text but with files."""
|
| 185 |
+
video_path = os.path.join(ROOT_DIR, "assets", "test_video.mp4")
|
| 186 |
+
|
| 187 |
+
message = {
|
| 188 |
+
"text": "", # Empty text
|
| 189 |
+
"files": [video_path]
|
| 190 |
+
}
|
| 191 |
+
|
| 192 |
+
# Add max_images parameter
|
| 193 |
+
result = process_user_input(message, 3)
|
| 194 |
+
|
| 195 |
+
# First item should be empty text
|
| 196 |
+
assert result[0]["type"] == "text"
|
| 197 |
+
assert result[0]["text"] == ""
|
| 198 |
+
|
| 199 |
+
# Rest should be video frames
|
| 200 |
+
assert len(result) > 1
|
| 201 |
+
|
| 202 |
+
|
| 203 |
+
def test_process_user_input_handles_empty_files_list():
|
| 204 |
+
"""Test that an empty files list is handled correctly."""
|
| 205 |
+
message = {
|
| 206 |
+
"text": "No files",
|
| 207 |
+
"files": []
|
| 208 |
+
}
|
| 209 |
+
|
| 210 |
+
# Add max_images parameter
|
| 211 |
+
result = process_user_input(message, 3)
|
| 212 |
+
assert len(result) == 1
|
| 213 |
+
assert result[0]["type"] == "text"
|
| 214 |
+
assert result[0]["text"] == "No files"
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
def test_process_user_input_max_images_effect():
|
| 218 |
+
"""Test that max_images parameter correctly limits the number of frames."""
|
| 219 |
+
video_path = os.path.join(ROOT_DIR, "assets", "test_video.mp4")
|
| 220 |
+
|
| 221 |
+
message = {
|
| 222 |
+
"text": "Video with few frames",
|
| 223 |
+
"files": [video_path]
|
| 224 |
+
}
|
| 225 |
+
|
| 226 |
+
result_few = process_user_input(message, 2)
|
| 227 |
+
result_many = process_user_input(message, 5)
|
| 228 |
+
|
| 229 |
+
# Count actual frames (each frame has a text and image entry)
|
| 230 |
+
frames_few = (len(result_few) - 1) // 2 # -1 for initial text message
|
| 231 |
+
frames_many = (len(result_many) - 1) // 2
|
| 232 |
+
|
| 233 |
+
# Should respect max_images parameter
|
| 234 |
+
assert frames_few <= 2
|
| 235 |
+
assert frames_many <= 5
|
| 236 |
+
assert frames_few < frames_many
|