feat: add search_by_appearance agent tool for clothing color search
- New Python script: clothing_color_search.py - New agent tool: search_by_appearance (red, blue, green, etc.) - Uses appearance.json person bboxes + HSV color analysis - Returns matched frames with confidence scores
This commit is contained in:
@@ -276,6 +276,15 @@ fn make_tools(pool: &sqlx::PgPool) -> Vec<ToolDef> {
|
||||
}),
|
||||
vec!["file_uuid", "node_id"],
|
||||
),
|
||||
function_calling::make_tool(
|
||||
"search_by_appearance",
|
||||
"根據衣服顏色搜尋影片中的人物。支援顏色:red, orange, yellow, green, cyan, blue, purple, white, black。",
|
||||
serde_json::json!({
|
||||
"file_uuid": {"type": "string", "description": "影片 UUID"},
|
||||
"color": {"type": "string", "description": "目標顏色: red, orange, yellow, green, cyan, blue, purple, white, black"}
|
||||
}),
|
||||
vec!["file_uuid", "color"],
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
@@ -310,6 +319,7 @@ async fn execute_tool(pool: &sqlx::PgPool, tool_call: &ToolCall) -> (String, Str
|
||||
"get_file_info" => tools::exec_get_file_info(pool, &args).await,
|
||||
"get_representative_frame" => tools::exec_get_representative_frame(pool, &args).await,
|
||||
"analyze_frame" => tools::exec_analyze_frame(pool, &args).await,
|
||||
"search_by_appearance" => tools::exec_search_by_appearance(pool, &args).await,
|
||||
_ => Err(format!("Unknown tool: {}", name)),
|
||||
};
|
||||
let content = match result {
|
||||
|
||||
@@ -1092,3 +1092,68 @@ pub async fn exec_tkg_node_detail(
|
||||
None => Err("Node not found".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Search for people by clothing color using appearance data
|
||||
pub async fn exec_search_by_appearance(pool: &sqlx::PgPool, args: &serde_json::Value) -> Result<String, String> {
|
||||
let file_uuid = args.get("file_uuid")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("file_uuid is required".to_string())?;
|
||||
let color = args.get("color")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("color is required (red, blue, green, yellow, orange, cyan, purple, white, black)".to_string())?;
|
||||
|
||||
let output_dir = std::env::var("MOMENTRY_OUTPUT_DIR")
|
||||
.unwrap_or_else(|_| "/Users/accusys/momentry/output".to_string());
|
||||
let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR")
|
||||
.unwrap_or_else(|_| "/Users/accusys/momentry_core/scripts".to_string());
|
||||
let script_path = format!("{}/clothing_color_search.py", scripts_dir);
|
||||
let appearance_path = format!("{}/{}.appearance.json", output_dir, file_uuid);
|
||||
let output_path = format!("{}/{}.color_search_{}.json", output_dir, file_uuid, color);
|
||||
|
||||
if !std::path::Path::new(&appearance_path).exists() {
|
||||
return Err(format!("appearance.json not found for file {}", file_uuid));
|
||||
}
|
||||
|
||||
// Get video path from videos table
|
||||
let videos_table = schema::table_name("videos");
|
||||
let video_path: Option<String> = sqlx::query_scalar(&format!(
|
||||
"SELECT file_path FROM {} WHERE file_uuid = $1", videos_table
|
||||
))
|
||||
.bind(file_uuid)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
let video_path = video_path.unwrap_or_default();
|
||||
if video_path.is_empty() {
|
||||
return Err("Video path not found".to_string());
|
||||
}
|
||||
|
||||
let executor = crate::core::processor::PythonExecutor::new()
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
executor.run(
|
||||
&script_path,
|
||||
&[
|
||||
"--file-uuid", file_uuid,
|
||||
"--color", color,
|
||||
"--video-path", &video_path,
|
||||
"--appearance-path", &appearance_path,
|
||||
"--output", &output_path,
|
||||
],
|
||||
None,
|
||||
"CLOTHING_COLOR_SEARCH",
|
||||
Some(std::time::Duration::from_secs(300)),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
// Read results
|
||||
if std::path::Path::new(&output_path).exists() {
|
||||
let content = std::fs::read_to_string(&output_path)
|
||||
.map_err(|e| e.to_string())?;
|
||||
Ok(content)
|
||||
} else {
|
||||
Err("Color search output not found".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user