5a8acece09
• Implemented robust model and scene reference validations in both the FastAPI backend and frontend to prevent empty, corrupt, or missing files from silently failing or producing duplicated background fallbacks. • Enhanced scenery prompt translation on the backend to dynamically map user-specified "Image" numbers e.g. Image 1, Image 2 to Qwen-specific "Picture" placeholders Picture 1, Picture 2, avoiding template mismatch desyncs. Changes • File & Path Validation: Added detailed validations to the /generate-scenery endpoint to ensure model_filename and scene_image paths exist on disk, are resolved without query parameters stripping ?t=..., are not directories, and are openable and verifiable as PIL images. • Identical Reference Checks: Introduced checks in both the backend and frontend submitGenerateScenery inside car.html to raise an HTTP 400 exception / display a toast warning if the model and scene references point to the same image. • Dynamic Prompt Mapping: Configured the backend to auto-translate occurrences of image 1/2/3 in custom scenery prompts to Picture 1/2/3, facilitating correct Qwen layout hint targeting. • Regression Unit Tests: Wrote comprehensive unit tests test_13b_scenery_validation_checks in test_regression_api.py verifying file checking, duplicate reference detection, and query parameter stripping. Verification • Executed test_regression_api.py simulated API testing with 100% success 19/19 tests passed. • Executed test_regression_qwen.py end-to-end model/hardware integration test suite successfully 2/2 tests passed.
1217 lines
53 KiB
Python
1217 lines
53 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import shutil
|
|
import unittest
|
|
from datetime import datetime as _dt
|
|
from PIL import Image
|
|
|
|
# Ensure tour-comfy is in import path so we can import the FastAPI app and database module
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from fastapi.testclient import TestClient
|
|
import database
|
|
from edit_api import app, _load_output_dir
|
|
|
|
class TestAPIRegression(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
# Determine output directory
|
|
cls.output_dir = _load_output_dir()
|
|
cls.client = TestClient(app)
|
|
|
|
# Unique mock identifiers to avoid conflicts
|
|
cls.test_ref_filename = "test_regression_ref_image_123.png"
|
|
cls.test_other_filename = "test_regression_other_image_123.png"
|
|
cls.group_id = "test_regression_group_123"
|
|
cls.test_ref_path = os.path.join(cls.output_dir, cls.test_ref_filename)
|
|
cls.test_other_path = os.path.join(cls.output_dir, cls.test_other_filename)
|
|
|
|
# Ensure any leftover artifacts from past failed runs are cleaned
|
|
cls._cleanup_database()
|
|
cls._cleanup_files()
|
|
|
|
# Create dummy test files (100x100 pixels, RGB format)
|
|
img = Image.new("RGB", (100, 100), color=(255, 0, 0))
|
|
img.save(cls.test_ref_path, "PNG")
|
|
|
|
img_other = Image.new("RGB", (100, 100), color=(0, 255, 0))
|
|
img_other.save(cls.test_other_path, "PNG")
|
|
|
|
# Mock metadata
|
|
cls.name = "Regression Test Character"
|
|
cls.tags = ["VISIBLE", "LIKE", "21+"]
|
|
cls.embedding = [0.1] * 1024
|
|
cls.clip_description = "A dummy regression test image description"
|
|
cls.prompt = "masterpiece, high quality"
|
|
cls.pose = "standing"
|
|
cls.group_name = "Regression Group"
|
|
cls.hidden = False
|
|
cls.has_background = True
|
|
cls.has_clothing = False
|
|
cls.is_source = True
|
|
cls.pose_description = "The model is standing and looking directly at the camera."
|
|
cls.pose_skeleton = '{"keypoints": [1, 2, 3]}'
|
|
|
|
# Insert reference image into the database
|
|
database.upsert_person(
|
|
cls.test_ref_filename,
|
|
filepath=cls.test_ref_path,
|
|
name=cls.name,
|
|
group_id=cls.group_id,
|
|
tags=cls.tags,
|
|
embedding=cls.embedding,
|
|
clip_description=cls.clip_description,
|
|
prompt=cls.prompt,
|
|
pose=cls.pose,
|
|
sort_order=0,
|
|
group_name=cls.group_name,
|
|
hidden=cls.hidden,
|
|
has_background=cls.has_background,
|
|
has_clothing=cls.has_clothing,
|
|
is_source=cls.is_source,
|
|
pose_description=cls.pose_description,
|
|
pose_skeleton=cls.pose_skeleton
|
|
)
|
|
|
|
# Insert second image in same group for reordering test
|
|
database.upsert_person(
|
|
cls.test_other_filename,
|
|
filepath=cls.test_other_path,
|
|
name=cls.name,
|
|
group_id=cls.group_id,
|
|
tags=cls.tags,
|
|
embedding=cls.embedding,
|
|
clip_description=cls.clip_description,
|
|
prompt=cls.prompt,
|
|
pose=cls.pose,
|
|
sort_order=1,
|
|
group_name=cls.group_name,
|
|
hidden=cls.hidden,
|
|
has_background=cls.has_background,
|
|
has_clothing=cls.has_clothing,
|
|
is_source=cls.is_source,
|
|
pose_description=cls.pose_description,
|
|
pose_skeleton=cls.pose_skeleton
|
|
)
|
|
|
|
# Track dynamically created files and database rows for cleanup
|
|
cls.created_files = [cls.test_ref_path, cls.test_other_path]
|
|
cls.created_db_rows = [cls.test_ref_filename, cls.test_other_filename]
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
# Cleanup
|
|
cls._cleanup_database()
|
|
cls._cleanup_files()
|
|
|
|
@classmethod
|
|
def _cleanup_database(cls):
|
|
# Safely delete any inserted test rows from the person database table
|
|
conn = database.get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
# Delete any filenames starting with test_regression or containing ts_crop/ts_pad
|
|
cur.execute("""
|
|
DELETE FROM person
|
|
WHERE filename LIKE 'test_regression_%%'
|
|
OR filename LIKE '%%_crop_test_regression_%%'
|
|
OR filename LIKE '%%_pad_test_regression_%%'
|
|
OR group_id = %s
|
|
""", (cls.group_id,))
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"Error cleaning up database: {e}")
|
|
conn.rollback()
|
|
finally:
|
|
cur.close()
|
|
database._put_db_connection(conn)
|
|
|
|
@classmethod
|
|
def _cleanup_files(cls):
|
|
# Clean up files matching test patterns
|
|
for f in os.listdir(cls.output_dir):
|
|
if "test_regression_" in f or "_crop_" in f or "_pad_" in f:
|
|
p = os.path.join(cls.output_dir, f)
|
|
try:
|
|
if os.path.exists(p):
|
|
os.remove(p)
|
|
except Exception as e:
|
|
print(f"Error removing file {p}: {e}")
|
|
|
|
def assertEmbeddingEqual(self, val1, val2):
|
|
if isinstance(val1, str):
|
|
val1 = [float(x) for x in val1.strip("[]").split(",")]
|
|
if isinstance(val2, str):
|
|
val2 = [float(x) for x in val2.strip("[]").split(",")]
|
|
self.assertEqual(len(val1), len(val2))
|
|
for x, y in zip(val1, val2):
|
|
self.assertAlmostEqual(x, y, places=4)
|
|
|
|
def test_01_duplicate_copies_all_pose_and_meta_details(self):
|
|
# Send duplicate request to FastAPI
|
|
response = self.client.post(f"/images/{self.test_ref_filename}/duplicate")
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertEqual(res_data["status"], "success")
|
|
|
|
new_filename = res_data["new_filename"]
|
|
self.created_db_rows.append(new_filename)
|
|
new_path = os.path.join(self.output_dir, os.path.basename(new_filename))
|
|
self.created_files.append(new_path)
|
|
|
|
# Assert duplicate file exists on disk
|
|
self.assertTrue(os.path.exists(new_path), f"Duplicated file {new_path} not found on disk")
|
|
|
|
# Assert all metadata has been accurately duplicated in DB
|
|
person = database.get_person(new_filename)
|
|
self.assertIsNotNone(person, "Duplicated database entry not found")
|
|
|
|
# Column mappings as in database.py get_person:
|
|
# SELECT name, group_id, tags, embedding, clip_description, filepath,
|
|
# prompt, pose, sort_order, group_name, hidden, has_background, source_refs,
|
|
# has_clothing, is_source, pose_description, pose_skeleton
|
|
self.assertEqual(person[0], self.name)
|
|
self.assertEqual(person[1], self.group_id)
|
|
|
|
# Tags (assert LIKE and 21+ are preserved)
|
|
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
|
self.assertIn("LIKE", tags_list)
|
|
self.assertIn("21+", tags_list)
|
|
self.assertIn("VISIBLE", tags_list)
|
|
|
|
# Verify 21+ is just a standard tag, not a safety blocker
|
|
self.assertIn("21+", tags_list)
|
|
|
|
self.assertEmbeddingEqual(person[3], self.embedding)
|
|
self.assertEqual(person[4], self.clip_description)
|
|
self.assertEqual(person[5], new_path)
|
|
self.assertEqual(person[6], self.prompt)
|
|
self.assertEqual(person[7], self.pose)
|
|
self.assertEqual(person[9], self.group_name)
|
|
self.assertEqual(person[10], self.hidden)
|
|
self.assertEqual(person[11], self.has_background)
|
|
|
|
# source_refs should refer to original
|
|
source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12]
|
|
self.assertIn(self.test_ref_filename, source_refs)
|
|
|
|
self.assertEqual(person[13], self.has_clothing)
|
|
self.assertEqual(person[14], self.is_source)
|
|
self.assertEqual(person[15], self.pose_description)
|
|
self.assertEqual(person[16], self.pose_skeleton)
|
|
|
|
def test_02_crop_as_copy_copies_all_pose_and_meta_details(self):
|
|
# Crop region: (10, 10, 90, 90), as_copy=True
|
|
req_payload = {
|
|
"x1": 10,
|
|
"y1": 10,
|
|
"x2": 90,
|
|
"y2": 90,
|
|
"as_copy": True
|
|
}
|
|
response = self.client.post(f"/images/{self.test_ref_filename}/crop", json=req_payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertEqual(res_data["status"], "success")
|
|
|
|
new_filename = res_data["new_filename"]
|
|
self.created_db_rows.append(new_filename)
|
|
new_path = os.path.join(self.output_dir, os.path.basename(new_filename))
|
|
self.created_files.append(new_path)
|
|
|
|
# Verify physical file existence and crop dimensions (should be 80x80)
|
|
self.assertTrue(os.path.exists(new_path), f"Cropped file {new_path} not found on disk")
|
|
cropped_img = Image.open(new_path)
|
|
self.assertEqual(cropped_img.size, (80, 80))
|
|
|
|
# Verify database entry has complete metadata
|
|
person = database.get_person(new_filename)
|
|
self.assertIsNotNone(person, "Cropped copy database entry not found")
|
|
|
|
self.assertEqual(person[0], self.name)
|
|
self.assertEqual(person[1], self.group_id)
|
|
|
|
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
|
self.assertIn("LIKE", tags_list)
|
|
self.assertIn("21+", tags_list)
|
|
self.assertIn("VISIBLE", tags_list)
|
|
|
|
self.assertEmbeddingEqual(person[3], self.embedding)
|
|
self.assertEqual(person[4], self.clip_description)
|
|
self.assertEqual(person[5], new_path)
|
|
self.assertEqual(person[6], self.prompt)
|
|
self.assertEqual(person[7], self.pose)
|
|
self.assertEqual(person[9], self.group_name)
|
|
self.assertEqual(person[10], self.hidden)
|
|
self.assertEqual(person[11], self.has_background)
|
|
|
|
source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12]
|
|
self.assertIn(self.test_ref_filename, source_refs)
|
|
|
|
self.assertEqual(person[13], self.has_clothing)
|
|
self.assertEqual(person[14], self.is_source)
|
|
self.assertEqual(person[15], self.pose_description)
|
|
self.assertEqual(person[16], self.pose_skeleton)
|
|
|
|
def test_03_crop_in_place_keeps_all_meta_information(self):
|
|
# Crop region: (20, 20, 80, 80), as_copy=False (in-place)
|
|
req_payload = {
|
|
"x1": 20,
|
|
"y1": 20,
|
|
"x2": 80,
|
|
"y2": 80,
|
|
"as_copy": False
|
|
}
|
|
|
|
# First verify original size is 100x100
|
|
orig_img = Image.open(self.test_ref_path)
|
|
self.assertEqual(orig_img.size, (100, 100))
|
|
|
|
response = self.client.post(f"/images/{self.test_ref_filename}/crop", json=req_payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertEqual(res_data["status"], "success")
|
|
self.assertEqual(res_data["new_filename"], self.test_ref_filename)
|
|
|
|
# Verify physical file size updated to 60x60
|
|
cropped_img = Image.open(self.test_ref_path)
|
|
self.assertEqual(cropped_img.size, (60, 60))
|
|
|
|
# Verify database entry has complete metadata untouched
|
|
person = database.get_person(self.test_ref_filename)
|
|
self.assertIsNotNone(person, "Database entry not found after in-place crop")
|
|
|
|
self.assertEqual(person[0], self.name)
|
|
self.assertEqual(person[1], self.group_id)
|
|
|
|
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
|
self.assertIn("LIKE", tags_list)
|
|
self.assertIn("21+", tags_list)
|
|
|
|
self.assertEmbeddingEqual(person[3], self.embedding)
|
|
self.assertEqual(person[4], self.clip_description)
|
|
self.assertEqual(person[6], self.prompt)
|
|
self.assertEqual(person[7], self.pose)
|
|
self.assertEqual(person[9], self.group_name)
|
|
self.assertEqual(person[10], self.hidden)
|
|
self.assertEqual(person[11], self.has_background)
|
|
self.assertEqual(person[13], self.has_clothing)
|
|
self.assertEqual(person[14], self.is_source)
|
|
self.assertEqual(person[15], self.pose_description)
|
|
self.assertEqual(person[16], self.pose_skeleton)
|
|
|
|
def test_04_pad_as_copy_copies_all_pose_and_meta_details(self):
|
|
# Pad canvas: expand top and bottom by 10 pixels, as_copy=True
|
|
req_payload = {
|
|
"top": 10,
|
|
"right": 0,
|
|
"bottom": 10,
|
|
"left": 0,
|
|
"as_copy": True,
|
|
"fill": "transparent",
|
|
"outpaint": False
|
|
}
|
|
# Original size at this point is 60x60 (due to test_03 crop)
|
|
response = self.client.post(f"/images/{self.test_ref_filename}/pad", json=req_payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertEqual(res_data["status"], "success")
|
|
|
|
new_filename = res_data["new_filename"]
|
|
self.created_db_rows.append(new_filename)
|
|
new_path = os.path.join(self.output_dir, os.path.basename(new_filename))
|
|
self.created_files.append(new_path)
|
|
|
|
# Verify physical file existence and pad dimensions (should be 60x80)
|
|
self.assertTrue(os.path.exists(new_path), f"Padded file {new_path} not found on disk")
|
|
padded_img = Image.open(new_path)
|
|
self.assertEqual(padded_img.size, (60, 80))
|
|
|
|
# Verify database entry has complete metadata
|
|
person = database.get_person(new_filename)
|
|
self.assertIsNotNone(person, "Padded copy database entry not found")
|
|
|
|
self.assertEqual(person[0], self.name)
|
|
self.assertEqual(person[1], self.group_id)
|
|
|
|
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
|
self.assertIn("LIKE", tags_list)
|
|
self.assertIn("21+", tags_list)
|
|
self.assertIn("VISIBLE", tags_list)
|
|
|
|
self.assertEmbeddingEqual(person[3], self.embedding)
|
|
self.assertEqual(person[4], self.clip_description)
|
|
self.assertEqual(person[5], new_path)
|
|
self.assertEqual(person[6], self.prompt)
|
|
self.assertEqual(person[7], self.pose)
|
|
self.assertEqual(person[9], self.group_name)
|
|
self.assertEqual(person[10], self.hidden)
|
|
self.assertEqual(person[11], self.has_background)
|
|
|
|
source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12]
|
|
self.assertIn(self.test_ref_filename, source_refs)
|
|
|
|
self.assertEqual(person[13], self.has_clothing)
|
|
self.assertEqual(person[14], self.is_source)
|
|
self.assertEqual(person[15], self.pose_description)
|
|
self.assertEqual(person[16], self.pose_skeleton)
|
|
|
|
def test_05_reorder_group_updates_sort_orders(self):
|
|
# Verify original order
|
|
order_rows_before = database.get_group_order(self.group_id)
|
|
filenames_before = [r[0] for r in order_rows_before]
|
|
self.assertIn(self.test_ref_filename, filenames_before)
|
|
self.assertIn(self.test_other_filename, filenames_before)
|
|
|
|
# Reverse the order of files and submit
|
|
new_order = [self.test_other_filename, self.test_ref_filename]
|
|
req_payload = {
|
|
"filenames": new_order
|
|
}
|
|
|
|
response = self.client.post(f"/groups/{self.group_id}/order", json=req_payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertEqual(res_data["group_id"], self.group_id)
|
|
self.assertEqual(res_data["filenames"], new_order)
|
|
|
|
# Get order from database and verify state reflects the updates
|
|
order_rows_after = database.get_group_order(self.group_id)
|
|
filenames_after = [r[0] for r in order_rows_after]
|
|
|
|
# Verify custom sort orders are explicitly 0, 1
|
|
self.assertEqual(filenames_after[:2], new_order)
|
|
|
|
person_other = database.get_person(self.test_other_filename)
|
|
person_ref = database.get_person(self.test_ref_filename)
|
|
|
|
# Column 8 in database.get_person is sort_order
|
|
self.assertEqual(person_other[8], 0)
|
|
self.assertEqual(person_ref[8], 1)
|
|
|
|
def test_06_list_images_with_bypass_static(self):
|
|
# Call with bypass_static=True
|
|
response = self.client.get("/images?archived=true&bypass_static=true")
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertIn("images", res_data)
|
|
|
|
# Find our reference image in the list
|
|
images = res_data["images"]
|
|
ref_img = next((img for img in images if img["filename"] == self.test_ref_filename), None)
|
|
self.assertIsNotNone(ref_img, "Test reference image not found in bypassed images list")
|
|
|
|
# Verify complete mapped metadata fields are correctly returned when bypassing static file
|
|
self.assertEqual(ref_img["name"], self.name)
|
|
self.assertEqual(ref_img["group_id"], self.group_id)
|
|
self.assertEqual(ref_img["prompt"], self.prompt)
|
|
self.assertEqual(ref_img["pose"], self.pose)
|
|
self.assertEqual(ref_img["group_name"], self.group_name)
|
|
self.assertEqual(ref_img["hidden"], self.hidden)
|
|
self.assertEqual(ref_img["has_background"], self.has_background)
|
|
self.assertEqual(ref_img["has_clothing"], self.has_clothing)
|
|
self.assertEqual(ref_img["is_source"], self.is_source)
|
|
self.assertEqual(ref_img["pose_description"], self.pose_description)
|
|
|
|
# Verify tags list
|
|
self.assertIn("LIKE", ref_img["tags"])
|
|
self.assertIn("21+", ref_img["tags"])
|
|
|
|
# Verify skeleton
|
|
self.assertEqual(ref_img["pose_skeleton"], self.pose_skeleton)
|
|
|
|
def test_07_concurrent_reordering_deadlock_resilience(self):
|
|
import threading
|
|
import random
|
|
import time
|
|
|
|
exceptions = []
|
|
|
|
def worker(thread_idx):
|
|
try:
|
|
for i in range(15):
|
|
# Alternate the order to trigger lock conflicts
|
|
if i % 2 == 0:
|
|
order = [self.test_other_filename, self.test_ref_filename]
|
|
else:
|
|
order = [self.test_ref_filename, self.test_other_filename]
|
|
|
|
database.set_group_order(self.group_id, order)
|
|
time.sleep(0.01)
|
|
except Exception as e:
|
|
exceptions.append((thread_idx, e))
|
|
|
|
threads = []
|
|
for j in range(5):
|
|
t = threading.Thread(target=worker, args=(j,))
|
|
threads.append(t)
|
|
t.start()
|
|
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# If any deadlock or serialization error was uncaught, exceptions won't be empty
|
|
if exceptions:
|
|
for tid, err in exceptions:
|
|
print(f"[test] Thread {tid} encountered error: {err}")
|
|
self.assertEqual(len(exceptions), 0, f"Encountered concurrent set_group_order errors: {exceptions}")
|
|
|
|
def test_08_deep_metadata_exposure_and_automatic_background_processing(self):
|
|
# Trigger the /images list
|
|
response = self.client.get("/images?archived=true&bypass_static=true")
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertIn("images", res_data)
|
|
|
|
# Verify that the new keys are present in each image dict
|
|
for img in res_data["images"]:
|
|
self.assertIn("people_count", img)
|
|
self.assertIn("anatomical_completeness", img)
|
|
self.assertIn("facial_direction", img)
|
|
self.assertIn("objects", img)
|
|
|
|
# Verify that manually triggering metadata backfill for our ref image returns expected structure
|
|
backfill_payload = {"filenames": [self.test_ref_filename]}
|
|
response = self.client.post("/images/backfill-metadata", json=backfill_payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
res_backfill = response.json()
|
|
self.assertEqual(res_backfill["status"], "completed")
|
|
self.assertEqual(res_backfill["total"], 1)
|
|
|
|
def test_09_upload_group_assignment(self):
|
|
from unittest.mock import patch
|
|
import io
|
|
|
|
# 1. Test skip_poses=True, group_id=None
|
|
# It should assign group_id to the base name of the uploaded file.
|
|
img_bytes = io.BytesIO()
|
|
Image.new("RGB", (100, 100), color=(0, 0, 255)).save(img_bytes, format="PNG")
|
|
img_bytes.seek(0)
|
|
|
|
response = self.client.post(
|
|
"/upload",
|
|
files={"image": ("test_upload_image.png", img_bytes, "image/png")},
|
|
data={"skip_poses": "true"}
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data = response.json()
|
|
self.assertEqual(res_data["status"], "added")
|
|
self.assertEqual(res_data["group_id"], "test_upload_image.png")
|
|
self.created_db_rows.append(res_data["filename"])
|
|
self.created_files.append(os.path.join(self.output_dir, res_data["filename"]))
|
|
|
|
# 2. Test skip_poses=False, group_id=None (e.g. CTRL+v New group with run poses)
|
|
# It should NOT collapse to "paste.png" or "test_upload_image.png", but generate a unique up_XXXX group ID!
|
|
img_bytes2 = io.BytesIO()
|
|
Image.new("RGB", (100, 100), color=(0, 255, 255)).save(img_bytes2, format="PNG")
|
|
img_bytes2.seek(0)
|
|
|
|
with patch("edit_api._process_upload") as mock_process:
|
|
response = self.client.post(
|
|
"/upload",
|
|
files={"image": ("test_upload_image2.png", img_bytes2, "image/png")},
|
|
data={"skip_poses": "false"}
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
res_data2 = response.json()
|
|
self.assertEqual(res_data2["status"], "processing")
|
|
# It should generate a unique ID starting with up_
|
|
self.assertTrue(res_data2["group_id"].startswith("up_"))
|
|
self.assertNotEqual(res_data2["group_id"], "test_upload_image2.png")
|
|
|
|
# Verify background task was called with the unique group_id
|
|
mock_process.assert_called_once()
|
|
args, kwargs = mock_process.call_args
|
|
# group_id is the 5th positional arg or a kwarg
|
|
called_gid = kwargs.get("group_id") or args[4]
|
|
self.assertTrue(called_gid.startswith("up_"))
|
|
|
|
self.created_files.append(os.path.join(self.output_dir, res_data2["filename"]))
|
|
|
|
def test_10_idle_backfill_non_image_filtering(self):
|
|
from unittest.mock import patch, MagicMock
|
|
import edit_api
|
|
|
|
# Reset failed backfill filenames set before testing
|
|
edit_api._failed_backfill_filenames.clear()
|
|
|
|
# Mock database.list_persons to return rows with:
|
|
# 1. A video file
|
|
# 2. A non-image file
|
|
# 3. An image that does exist on disk but we want to backfill
|
|
# 4. A normal image where people_count is already present
|
|
mock_persons = [
|
|
("video.mp4", "Video", "g1", None, "", None, 0, "", False, False, None, False, "video", None, False, False, [], None, None, None, None, None, []),
|
|
("notes.txt", "Text", "g1", None, "", None, 1, "", False, False, None, False, "image", None, False, False, [], None, None, None, None, None, []),
|
|
("processed_image.png", "Image1", "g1", None, "", None, 2, "", False, False, None, False, "image", None, False, False, [], None, None, 1, "Full", "Front", []),
|
|
("test_ref_image_123.png", "Image2", "g1", None, "", None, 3, "", False, False, None, False, "image", None, False, False, [], None, None, None, None, None, []),
|
|
]
|
|
|
|
# First, test the fast-fail guard inside _process_image_for_metadata directly (not mocked yet)
|
|
res_mp4 = edit_api._process_image_for_metadata("video.mp4")
|
|
self.assertIsNone(res_mp4)
|
|
|
|
res_txt = edit_api._process_image_for_metadata("notes.txt")
|
|
self.assertIsNone(res_txt)
|
|
|
|
# Let's mock os.path.exists to return True for any path we check
|
|
with patch("os.path.exists", return_value=True), \
|
|
patch("database.list_persons", return_value=mock_persons), \
|
|
patch("edit_api._process_image_for_metadata") as mock_process:
|
|
|
|
# Second, let's replicate/test the candidate selection loop inside edit_api._idle_turntable_daemon
|
|
output_dir = "/mnt/zim/tour-comfy/output"
|
|
_failed_backfill_filenames = edit_api._failed_backfill_filenames
|
|
|
|
persons = mock_persons
|
|
legacy_candidate = None
|
|
for row in persons:
|
|
fname = row[0]
|
|
content_type = row[12] if len(row) > 12 else None
|
|
people_count = row[19] if len(row) > 19 else None
|
|
if fname.startswith("_turntable/"):
|
|
continue
|
|
if content_type == 'video':
|
|
continue
|
|
if not fname.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
|
|
continue
|
|
fpath = os.path.join(output_dir, fname)
|
|
if not os.path.exists(fpath):
|
|
continue
|
|
if people_count is None and fname not in _failed_backfill_filenames:
|
|
legacy_candidate = fname
|
|
break
|
|
|
|
self.assertEqual(legacy_candidate, "test_ref_image_123.png")
|
|
|
|
# Let's test failed backfill insertion if _process_image_for_metadata returns None
|
|
mock_process.return_value = None
|
|
|
|
# Simulating the idle turntable loop processing the candidate
|
|
if legacy_candidate:
|
|
try:
|
|
res = edit_api._process_image_for_metadata(legacy_candidate)
|
|
if res is None:
|
|
_failed_backfill_filenames.add(legacy_candidate)
|
|
except Exception as ex:
|
|
_failed_backfill_filenames.add(legacy_candidate)
|
|
|
|
self.assertIn("test_ref_image_123.png", _failed_backfill_filenames)
|
|
|
|
def test_11_video_listing_refresh_and_trim_synchronization(self):
|
|
from unittest.mock import patch
|
|
import edit_api
|
|
|
|
# Mock wireframe folder and output directories
|
|
with patch("edit_api._load_wireframe_dir", return_value="/tmp/test_wireframe"), \
|
|
patch("edit_api._load_output_dir", return_value=self.output_dir), \
|
|
patch("os.path.isdir", return_value=True), \
|
|
patch("os.listdir", return_value=["dance.mp4", "dance_10s-20s.mp4", "notes.txt"]):
|
|
|
|
# Verify that calling /videos?refresh=true triggers grouping correctly
|
|
response = self.client.get("/videos?refresh=true")
|
|
self.assertEqual(response.status_code, 200)
|
|
data = response.json()
|
|
|
|
# Grouping checks
|
|
self.assertIn("groups", data)
|
|
groups = data["groups"]
|
|
# Find group with stem "dance"
|
|
dance_grp = next((g for g in groups if g["stem"] == "dance"), None)
|
|
self.assertIsNotNone(dance_grp)
|
|
self.assertEqual(dance_grp["video"], "dance.mp4")
|
|
self.assertIn("dance_10s-20s.mp4", dance_grp["clips"])
|
|
|
|
def test_12_standalone_scenery_sorting_and_existence_check(self):
|
|
from unittest.mock import patch
|
|
import edit_api
|
|
import database
|
|
|
|
# Insert some mock scenery generation records so we can test last_used sorting
|
|
conn = database.get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
INSERT INTO person (filename, group_id, source_refs)
|
|
VALUES ('20260629_130000_sc_test.png', 'g1', '["test_person.png", "wireframe:used_scenery_2.png"]')
|
|
""")
|
|
cur.execute("""
|
|
INSERT INTO person (filename, group_id, source_refs)
|
|
VALUES ('20260629_130500_sc_test.png', 'g1', '["test_person.png", "wireframe:used_scenery_1.png"]')
|
|
""")
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
database._put_db_connection(conn)
|
|
|
|
try:
|
|
# Mock wireframe folder and output directories
|
|
with patch("edit_api._load_wireframe_dir", return_value="/tmp/test_wireframe"), \
|
|
patch("edit_api._load_output_dir", return_value=self.output_dir), \
|
|
patch("os.path.isdir", return_value=True), \
|
|
patch("os.path.exists", return_value=True), \
|
|
patch("os.path.getmtime", return_value=12345.0), \
|
|
patch("os.listdir", return_value=["used_scenery_1.png", "used_scenery_2.png", "unused_scenery_3.png"]):
|
|
|
|
response = self.client.get("/videos?refresh=true")
|
|
self.assertEqual(response.status_code, 200)
|
|
data = response.json()
|
|
|
|
# Check standalone images exist in the response
|
|
self.assertIn("standalone_images", data)
|
|
standalone = data["standalone_images"]
|
|
|
|
# Should contain our wireframes sorted with last used on top:
|
|
# 'used_scenery_1.png' was used in the most recent scenery generation (at 13:05:00)
|
|
# 'used_scenery_2.png' was used in the older scenery generation (at 13:00:00)
|
|
# 'unused_scenery_3.png' was not used
|
|
self.assertEqual(standalone[0], "used_scenery_1.png")
|
|
self.assertEqual(standalone[1], "used_scenery_2.png")
|
|
self.assertEqual(standalone[2], "unused_scenery_3.png")
|
|
|
|
# Also verify that /scenery/library filters out missing files
|
|
# Let's mock _get_cached_file_meta: '20260629_130500_sc_test.png' exists, but '20260629_130000_sc_test.png' is missing
|
|
def mock_cached_meta(fname, out_dir):
|
|
if fname == "20260629_130500_sc_test.png":
|
|
return True, 12345.0
|
|
return False, 0.0
|
|
|
|
with patch("edit_api._get_cached_file_meta", side_effect=mock_cached_meta):
|
|
response = self.client.get("/scenery/library")
|
|
self.assertEqual(response.status_code, 200)
|
|
lib_data = response.json()
|
|
|
|
# Verify that only the existing scenery item '20260629_130500_sc_test.png' is returned
|
|
all_filenames = []
|
|
for g in lib_data["groups"]:
|
|
for item in g["items"]:
|
|
all_filenames.append(item["filename"])
|
|
for item in lib_data["ungrouped"]:
|
|
all_filenames.append(item["filename"])
|
|
|
|
self.assertIn("20260629_130500_sc_test.png", all_filenames)
|
|
self.assertNotIn("20260629_130000_sc_test.png", all_filenames)
|
|
|
|
finally:
|
|
# Clean up test DB rows
|
|
conn = database.get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM person WHERE filename IN ('20260629_130000_sc_test.png', '20260629_130500_sc_test.png')")
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
database._put_db_connection(conn)
|
|
|
|
def test_13_scenery_metadata_extraction_and_refs(self):
|
|
from unittest.mock import patch, MagicMock
|
|
import edit_api
|
|
import database
|
|
import json
|
|
import os
|
|
from PIL import Image
|
|
|
|
# Write dummy model image on disk
|
|
model_filename = "model_image_123.png"
|
|
model_path = os.path.join(self.output_dir, model_filename)
|
|
Image.new("RGB", (10, 10), color="red").save(model_path)
|
|
|
|
# Set up a test person record in DB for the model
|
|
conn = database.get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
INSERT INTO person (filename, group_id)
|
|
VALUES (%s, 'model_group_123')
|
|
""", (model_filename,))
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
database._put_db_connection(conn)
|
|
|
|
try:
|
|
# Let's mock _run_pipeline, generation of embeddings, etc.
|
|
# We want to verify that _scenery_worker correctly structures 'source_refs'
|
|
# and triggers metadata extraction.
|
|
mock_png = b"dummy_png_bytes"
|
|
edit_api.jobs["test_scenery_job"] = {"status": "running"}
|
|
|
|
with patch("edit_api._run_pipeline", return_value=mock_png), \
|
|
patch("edit_api._load_output_dir", return_value=self.output_dir), \
|
|
patch("edit_api.embeddings.generate_embedding", return_value=None), \
|
|
patch("edit_api._metadata_executor.submit") as mock_submit:
|
|
|
|
# Run the scenery worker synchronously for the test
|
|
dummy_scene = Image.new("RGB", (100, 100), color="blue")
|
|
|
|
edit_api._scenery_worker(
|
|
job_id="test_scenery_job",
|
|
model_filename=model_filename,
|
|
scene_pil=dummy_scene,
|
|
prompt="test scenery prompt",
|
|
seed=42,
|
|
extra_pils=[],
|
|
scene_video="test_video.mp4",
|
|
scene_image="test_image.png",
|
|
extra_filename="test_extra.png"
|
|
)
|
|
|
|
# Find the newly created scenery filename
|
|
scenery_filename = edit_api.jobs["test_scenery_job"]["output"]
|
|
self.assertIsNotNone(scenery_filename)
|
|
|
|
# Check database record to verify refs are correct
|
|
person = database.get_person(scenery_filename)
|
|
self.assertIsNotNone(person)
|
|
self.assertEqual(person[1], "model_group_123") # Correctly pulled group_id
|
|
|
|
refs = json.loads(person[12]) # source_refs is at index 12 in tuple or get_person response
|
|
self.assertIn(model_filename, refs)
|
|
self.assertIn("video:test_video.mp4", refs)
|
|
self.assertIn("wireframe:test_image.png", refs)
|
|
self.assertIn("test_extra.png", refs)
|
|
|
|
# Verify that _metadata_executor.submit was called to trigger background metadata extraction
|
|
mock_submit.assert_called_once_with(edit_api._process_image_for_metadata, scenery_filename)
|
|
|
|
finally:
|
|
# Clean up all created files and database entries
|
|
if os.path.exists(model_path):
|
|
os.remove(model_path)
|
|
conn = database.get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM person WHERE filename = 'model_image_123.png'")
|
|
cur.execute("DELETE FROM person WHERE filename LIKE '%_sc_model_image_123.png'")
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
database._put_db_connection(conn)
|
|
|
|
# Clean up generated scenery file on disk
|
|
try:
|
|
scenery_filename = edit_api.jobs.get("test_scenery_job", {}).get("output")
|
|
if scenery_filename:
|
|
out_path = os.path.join(self.output_dir, scenery_filename)
|
|
if os.path.exists(out_path):
|
|
os.remove(out_path)
|
|
except Exception:
|
|
pass
|
|
|
|
def test_13b_scenery_validation_checks(self):
|
|
from unittest.mock import patch
|
|
import edit_api
|
|
from fastapi import HTTPException
|
|
|
|
# Test 1: model image is completely missing -> raise 404
|
|
req_missing = edit_api.SceneryRequest(
|
|
model_filename="non_existent_model_file_999.png",
|
|
scene_image="some_scene.png",
|
|
prompt="replace person from image 1 with the person from Image 2"
|
|
)
|
|
with self.assertRaises(HTTPException) as ctx:
|
|
edit_api.generate_scenery(req_missing)
|
|
self.assertEqual(ctx.exception.status_code, 404)
|
|
self.assertIn("Model image not found", ctx.exception.detail)
|
|
|
|
# Write dummy model image on disk for subsequent tests
|
|
model_filename = "test_model_validation_123.png"
|
|
model_path = os.path.join(self.output_dir, model_filename)
|
|
Image.new("RGB", (10, 10), color="red").save(model_path)
|
|
|
|
try:
|
|
# Test 2: model and scene images are the same file -> raise 400
|
|
req_identical = edit_api.SceneryRequest(
|
|
model_filename=model_filename,
|
|
scene_image=model_filename,
|
|
prompt="replace person from image 1 with the person from Image 2"
|
|
)
|
|
with self.assertRaises(HTTPException) as ctx:
|
|
edit_api.generate_scenery(req_identical)
|
|
self.assertEqual(ctx.exception.status_code, 400)
|
|
self.assertIn("cannot be the same image", ctx.exception.detail)
|
|
|
|
# Test 3: prompt refinement "image 1/2" -> "Picture 1/2"
|
|
# We mock the thread start and save_db_prompt to verify the processed prompt
|
|
saved_prompt = []
|
|
def mock_save_db_prompt(ptype, prompt_text, meta):
|
|
saved_prompt.append(prompt_text)
|
|
|
|
dummy_scene_filename = "dummy_scene_validation_123.png"
|
|
dummy_scene_path = os.path.join(self.output_dir, dummy_scene_filename)
|
|
Image.new("RGB", (10, 10), color="blue").save(dummy_scene_path)
|
|
|
|
try:
|
|
req_refine = edit_api.SceneryRequest(
|
|
model_filename=model_filename,
|
|
scene_image=dummy_scene_filename,
|
|
prompt="replace person from image 1 with the person from Image 2"
|
|
)
|
|
|
|
with patch("edit_api._load_wireframe_dir", return_value=self.output_dir), \
|
|
patch("edit_api.database.save_db_prompt", side_effect=mock_save_db_prompt), \
|
|
patch("threading.Thread.start") as mock_thread_start:
|
|
|
|
edit_api.generate_scenery(req_refine)
|
|
|
|
self.assertEqual(len(saved_prompt), 1)
|
|
self.assertEqual(saved_prompt[0], "replace person from Picture 1 with the person from Picture 2")
|
|
|
|
finally:
|
|
if os.path.exists(dummy_scene_path):
|
|
os.remove(dummy_scene_path)
|
|
|
|
finally:
|
|
if os.path.exists(model_path):
|
|
os.remove(model_path)
|
|
|
|
def test_13c_scenery_generation_insufficient_references_checks(self):
|
|
import edit_api
|
|
from fastapi import HTTPException
|
|
from PIL import Image
|
|
|
|
# Write dummy model image on disk
|
|
model_filename = "test_scenery_insufficient_model.png"
|
|
model_path = os.path.join(self.output_dir, model_filename)
|
|
Image.new("RGB", (10, 10), color="red").save(model_path)
|
|
|
|
try:
|
|
# Test 1: both scene_image and scene_video are missing -> raise HTTPException 400
|
|
req_no_bg = edit_api.SceneryRequest(
|
|
model_filename=model_filename,
|
|
prompt="some prompt"
|
|
)
|
|
with self.assertRaises(HTTPException) as ctx:
|
|
edit_api.generate_scenery(req_no_bg)
|
|
self.assertEqual(ctx.exception.status_code, 400)
|
|
self.assertIn("requires a background scene reference", ctx.exception.detail)
|
|
|
|
# Test 2: inside _scenery_worker, if len(refs) < 2, it raises ValueError and sets status to error
|
|
edit_api.jobs["test_scenery_insufficient_job"] = {"status": "running"}
|
|
dummy_scene = Image.new("RGB", (10, 10), color="blue")
|
|
edit_api._scenery_worker(
|
|
job_id="test_scenery_insufficient_job",
|
|
model_filename=model_filename,
|
|
scene_pil=dummy_scene,
|
|
prompt="test scenery prompt",
|
|
seed=42,
|
|
extra_pils=[],
|
|
scene_video=None, # explicitly None to trigger the < 2 refs check
|
|
scene_image=None, # explicitly None to trigger the < 2 refs check
|
|
extra_filename=None
|
|
)
|
|
job = edit_api.jobs["test_scenery_insufficient_job"]
|
|
self.assertEqual(job["status"], "error")
|
|
self.assertIn("expected at least 2 source references", job["error"])
|
|
|
|
finally:
|
|
if os.path.exists(model_path):
|
|
os.remove(model_path)
|
|
|
|
def test_14_static_serialization_of_deep_metadata(self):
|
|
import edit_api
|
|
import database
|
|
import json
|
|
import os
|
|
from PIL import Image
|
|
|
|
# Create dummy image and person record
|
|
fn = "test_deep_meta_serialization_123.png"
|
|
fpath = os.path.join(self.output_dir, fn)
|
|
Image.new("RGB", (10, 10)).save(fpath)
|
|
|
|
conn = database.get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
INSERT INTO person (filename, filepath, group_id, people_count, anatomical_completeness, facial_direction, objects)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
|
""", (fn, fpath, "test_group", 2, False, "front", json.dumps([{"tag": "hat", "score": 0.9}])))
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
database._put_db_connection(conn)
|
|
|
|
try:
|
|
# Trigger write static images.json
|
|
edit_api._write_all_static()
|
|
|
|
# Read static file and check that the deep metadata properties are serialized correctly
|
|
static_file = os.path.join(self.output_dir, "_data", "images.json")
|
|
self.assertTrue(os.path.exists(static_file))
|
|
with open(static_file, "r") as f:
|
|
data = json.load(f)
|
|
|
|
imgs = data.get("images", [])
|
|
matched = [x for x in imgs if x["filename"] == fn]
|
|
self.assertEqual(len(matched), 1)
|
|
item = matched[0]
|
|
|
|
# Assert that the serialized dictionary has all of the deep metadata fields
|
|
self.assertEqual(item.get("people_count"), 2)
|
|
self.assertEqual(item.get("anatomical_completeness"), False)
|
|
self.assertEqual(item.get("facial_direction"), "front")
|
|
self.assertEqual(item.get("objects"), [{"tag": "hat", "score": 0.9}])
|
|
|
|
finally:
|
|
# Clean up
|
|
if os.path.exists(fpath):
|
|
os.remove(fpath)
|
|
conn = database.get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM person WHERE filename = %s", (fn,))
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
database._put_db_connection(conn)
|
|
# Re-generate static to clean up the dummy entry
|
|
edit_api._write_all_static()
|
|
|
|
def test_15_anatomical_completeness_refinement(self):
|
|
import edit_api
|
|
|
|
# 17 keypoints matching COCO-17 format: [x, y, score]
|
|
# Nose (0), shoulders (5,6), elbows (7,8), wrists (9,10), hips (11,12), knees (13,14), ankles (15,16)
|
|
# All of them are visible with score >= 0.3, and well within bounds [0.1, 0.9]
|
|
complete_kpts = [
|
|
[50, 10, 0.9], # Nose (0)
|
|
[48, 11, 0.9], # L_Eye (1)
|
|
[52, 11, 0.9], # R_Eye (2)
|
|
[46, 12, 0.9], # L_Ear (3)
|
|
[54, 12, 0.9], # R_Ear (4)
|
|
[40, 20, 0.9], # L_Shoulder (5)
|
|
[60, 20, 0.9], # R_Shoulder (6)
|
|
[35, 35, 0.9], # L_Elbow (7)
|
|
[65, 35, 0.9], # R_Elbow (8)
|
|
[30, 50, 0.9], # L_Wrist (9)
|
|
[70, 50, 0.9], # R_Wrist (10)
|
|
[42, 55, 0.9], # L_Hip (11)
|
|
[58, 55, 0.9], # R_Hip (12)
|
|
[40, 75, 0.9], # L_Knee (13)
|
|
[60, 75, 0.9], # R_Knee (14)
|
|
[40, 90, 0.9], # L_Ankle (15)
|
|
[60, 90, 0.9], # R_Ankle (16)
|
|
]
|
|
|
|
# Test complete pose
|
|
res1 = edit_api._detect_anatomical_completeness(complete_kpts, 100, 100)
|
|
self.assertTrue(res1)
|
|
|
|
# Test partial pose: missing knees and ankles
|
|
partial_kpts = [list(kp) for kp in complete_kpts]
|
|
for idx in [13, 14, 15, 16]:
|
|
partial_kpts[idx][2] = 0.1 # Hide them
|
|
res2 = edit_api._detect_anatomical_completeness(partial_kpts, 100, 100)
|
|
self.assertFalse(res2)
|
|
|
|
# Test cropped pose: ankles too close to bottom edge (y=99 in height=100)
|
|
cropped_bottom_kpts = [list(kp) for kp in complete_kpts]
|
|
cropped_bottom_kpts[15][1] = 99
|
|
cropped_bottom_kpts[16][1] = 99
|
|
res3 = edit_api._detect_anatomical_completeness(cropped_bottom_kpts, 100, 100)
|
|
self.assertFalse(res3)
|
|
|
|
def test_16_object_bbox_estimation(self):
|
|
import edit_api
|
|
|
|
# Define some realistic keypoints
|
|
kpts = [
|
|
[50, 10, 0.9], # Nose (0)
|
|
[48, 11, 0.9], # L_Eye (1)
|
|
[52, 11, 0.9], # R_Eye (2)
|
|
[46, 12, 0.9], # L_Ear (3)
|
|
[54, 12, 0.9], # R_Ear (4)
|
|
[40, 20, 0.9], # L_Shoulder (5)
|
|
[60, 20, 0.9], # R_Shoulder (6)
|
|
[35, 35, 0.9], # L_Elbow (7)
|
|
[65, 35, 0.9], # R_Elbow (8)
|
|
[30, 50, 0.9], # L_Wrist (9)
|
|
[70, 50, 0.9], # R_Wrist (10)
|
|
[42, 55, 0.9], # L_Hip (11)
|
|
[58, 55, 0.9], # R_Hip (12)
|
|
[40, 75, 0.9], # L_Knee (13)
|
|
[60, 75, 0.9], # R_Knee (14)
|
|
[40, 90, 0.9], # L_Ankle (15)
|
|
[60, 90, 0.9], # R_Ankle (16)
|
|
]
|
|
|
|
# Test head / hair tag bbox estimation
|
|
bbox_hair = edit_api._estimate_bbox_for_tag("long_hair", kpts, 100, 100)
|
|
self.assertIsNotNone(bbox_hair)
|
|
self.assertEqual(len(bbox_hair), 4)
|
|
# Bounding box should surround the head/face region, which is y-centered around 10-12.
|
|
# So upper y (bbox_hair[1]) should be relatively small (close to 0/top)
|
|
self.assertLess(bbox_hair[1], 15)
|
|
|
|
# Test chest tag bbox estimation
|
|
bbox_breasts = edit_api._estimate_bbox_for_tag("breasts", kpts, 100, 100)
|
|
self.assertIsNotNone(bbox_breasts)
|
|
self.assertEqual(len(bbox_breasts), 4)
|
|
# Should be below shoulders (y=20) and above hips (y=55). Center around 30.
|
|
self.assertGreater(bbox_breasts[1], 15)
|
|
self.assertLess(bbox_breasts[3], 55)
|
|
|
|
# Test stomach tag bbox estimation
|
|
bbox_navel = edit_api._estimate_bbox_for_tag("navel", kpts, 100, 100)
|
|
self.assertIsNotNone(bbox_navel)
|
|
self.assertEqual(len(bbox_navel), 4)
|
|
# Should be around the stomach area, i.e. between 35 and 55.
|
|
self.assertGreater(bbox_navel[1], 25)
|
|
|
|
# Test arms tag bbox estimation
|
|
bbox_arms = edit_api._estimate_bbox_for_tag("sleeves", kpts, 100, 100)
|
|
self.assertIsNotNone(bbox_arms)
|
|
self.assertEqual(len(bbox_arms), 4)
|
|
|
|
# Test legs tag bbox estimation
|
|
bbox_legs = edit_api._estimate_bbox_for_tag("thighs", kpts, 100, 100)
|
|
self.assertIsNotNone(bbox_legs)
|
|
self.assertEqual(len(bbox_legs), 4)
|
|
|
|
def test_17_gaze_and_pose_description_enhancements(self):
|
|
import edit_api
|
|
|
|
# Test 1: Standard upright front facing pose with direct gaze
|
|
standing_kpts = [
|
|
[50, 12.5, 0.9], # Nose (0)
|
|
[52, 11, 0.9], # L_Eye (1)
|
|
[48, 11, 0.9], # R_Eye (2)
|
|
[54, 12, 0.9], # L_Ear (3)
|
|
[46, 12, 0.9], # R_Ear (4)
|
|
[60, 20, 0.9], # L_Shoulder (5)
|
|
[40, 20, 0.9], # R_Shoulder (6)
|
|
[65, 35, 0.9], # L_Elbow (7)
|
|
[35, 35, 0.9], # R_Elbow (8)
|
|
[70, 50, 0.9], # L_Wrist (9)
|
|
[30, 50, 0.9], # R_Wrist (10)
|
|
[58, 55, 0.9], # L_Hip (11)
|
|
[42, 55, 0.9], # R_Hip (12)
|
|
[60, 75, 0.9], # L_Knee (13)
|
|
[40, 75, 0.9], # R_Knee (14)
|
|
[60, 90, 0.9], # L_Ankle (15)
|
|
[40, 90, 0.9], # R_Ankle (16)
|
|
]
|
|
|
|
pose_desc = edit_api._describe_pose(standing_kpts)
|
|
self.assertIn("standing", pose_desc)
|
|
self.assertIn("facing forward", pose_desc)
|
|
|
|
gaze = edit_api._detect_facial_direction(standing_kpts)
|
|
self.assertEqual(gaze, "looking forward")
|
|
|
|
# Test 2: Gaze look left and up
|
|
left_up_kpts = [list(kp) for kp in standing_kpts]
|
|
# To look left: nose (0) moves to the left (smaller X than ear midpoint 50)
|
|
left_up_kpts[0][0] = 45 # Nose X
|
|
# To look up: nose (0) moves up (y-level close to eye level y=11)
|
|
left_up_kpts[0][1] = 11.2 # Nose Y
|
|
|
|
gaze_left_up = edit_api._detect_facial_direction(left_up_kpts)
|
|
self.assertEqual(gaze_left_up, "looking left and up")
|
|
|
|
# Test 3: Lying down / reclining pose
|
|
lying_kpts = [list(kp) for kp in standing_kpts]
|
|
# In lying down, nose is at y=10, hips are at y=12 (very horizontal)
|
|
# Head X = 10, Hip X = 80
|
|
lying_kpts[0][0] = 10 # Head X
|
|
lying_kpts[0][1] = 10 # Head Y
|
|
lying_kpts[11][0] = 80 # Hip X
|
|
lying_kpts[11][1] = 12 # Hip Y
|
|
lying_kpts[12][0] = 80
|
|
lying_kpts[12][1] = 12
|
|
|
|
pose_desc_lying = edit_api._describe_pose(lying_kpts)
|
|
self.assertIn("lying down", pose_desc_lying)
|
|
|
|
def test_18_designer_enhancements(self):
|
|
import unittest.mock as mock
|
|
import edit_api
|
|
|
|
# 1. Test Deduplication of Pose Names
|
|
# We will mock the external requests.post LLM call to return some mock poses,
|
|
# one of which collides with an existing pose name in poses.md (e.g., "The Clasp" or "standing")
|
|
mock_raw_response = (
|
|
"# standing\n"
|
|
"This is a standing pose description.\n"
|
|
"Perfect anatomy, photorealistic.\n\n"
|
|
"# Custom Pose\n"
|
|
"This is a custom pose description.\n"
|
|
"Anatomically precise, photorealistic.\n"
|
|
)
|
|
|
|
mock_response = mock.Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": mock_raw_response
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
with mock.patch("requests.post", return_value=mock_response) as mock_post:
|
|
# Let's override _load_poses to return {"standing": "existing text"} to force a name collision
|
|
with mock.patch("edit_api._load_poses", return_value={"standing": {"text": "existing text", "beta": False}}):
|
|
payload = {
|
|
"n": 2,
|
|
"context": "standing and custom pose instructions",
|
|
"filename": None,
|
|
"beta": False,
|
|
"messages": None
|
|
}
|
|
response = self.client.post("/designer/generate", json=payload)
|
|
self.assertEqual(response.status_code, 200)
|
|
data = response.json()
|
|
self.assertEqual(data["status"], "success")
|
|
poses = data["poses"]
|
|
# The duplicate "standing" name should be renamed to "standing 2" instead of skipped!
|
|
self.assertIn("standing 2", poses)
|
|
self.assertIn("Custom Pose", poses)
|
|
self.assertEqual(poses["standing 2"], "This is a standing pose description. Perfect anatomy, photorealistic.")
|
|
self.assertEqual(poses["Custom Pose"], "This is a custom pose description. Anatomically precise, photorealistic.")
|
|
|
|
# 2. Test Multi-Turn Conversation/History Payload Building
|
|
with mock.patch("requests.post", return_value=mock_response) as mock_post:
|
|
with mock.patch("edit_api._load_poses", return_value={}):
|
|
# Send history messages in the request
|
|
history_messages = [
|
|
{"role": "user", "content": "Initial prompt"},
|
|
{"role": "assistant", "content": "# Some Pose\nbody text"}
|
|
]
|
|
payload_with_history = {
|
|
"n": 1,
|
|
"context": "make them sit down",
|
|
"filename": None,
|
|
"beta": False,
|
|
"messages": history_messages
|
|
}
|
|
response = self.client.post("/designer/generate", json=payload_with_history)
|
|
self.assertEqual(response.status_code, 200)
|
|
|
|
# Check that requests.post was called with the correct messages payload
|
|
self.assertTrue(mock_post.called)
|
|
call_args = mock_post.call_args
|
|
called_json = call_args[1]["json"]
|
|
called_messages = called_json["messages"]
|
|
|
|
# The payload should contain:
|
|
# 1. System message (DESIGNER_SYSTEM)
|
|
# 2. Initial user prompt
|
|
# 3. Assistant response
|
|
# 4. New user follow-up prompt incorporating "make them sit down"
|
|
self.assertEqual(len(called_messages), 4)
|
|
self.assertEqual(called_messages[0]["role"], "system")
|
|
self.assertEqual(called_messages[1]["role"], "user")
|
|
self.assertEqual(called_messages[1]["content"], "Initial prompt")
|
|
self.assertEqual(called_messages[2]["role"], "assistant")
|
|
self.assertEqual(called_messages[2]["content"], "# Some Pose\nbody text")
|
|
self.assertEqual(called_messages[3]["role"], "user")
|
|
self.assertIn("make them sit down", called_messages[3]["content"])
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|