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:
216
scripts/clothing_color_search.py
Normal file
216
scripts/clothing_color_search.py
Normal file
@@ -0,0 +1,216 @@
|
|||||||
|
#!/opt/homebrew/bin/python3.11
|
||||||
|
"""
|
||||||
|
Clothing Color Search - Find people wearing specific colors
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python3 clothing_color_search.py --file-uuid UUID --color red --output output.json
|
||||||
|
|
||||||
|
Color matching uses HSV hue ranges:
|
||||||
|
red: 0-15, 165-180
|
||||||
|
orange: 15-35
|
||||||
|
yellow: 35-50
|
||||||
|
green: 50-85
|
||||||
|
cyan: 85-105
|
||||||
|
blue: 105-140
|
||||||
|
purple: 140-165
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
import json
|
||||||
|
import argparse
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
COLOR_RANGES = {
|
||||||
|
"red": [(0, 40), (165, 180)],
|
||||||
|
"orange": [(15, 35)],
|
||||||
|
"yellow": [(35, 50)],
|
||||||
|
"green": [(50, 85)],
|
||||||
|
"cyan": [(85, 105)],
|
||||||
|
"blue": [(105, 140)],
|
||||||
|
"purple": [(140, 165)],
|
||||||
|
"white": [(0, 180, 0, 40, 200, 255)], # (h_min, h_max, s_min, s_max, v_min, v_max)
|
||||||
|
"black": [(0, 180, 0, 255, 0, 50)],
|
||||||
|
}
|
||||||
|
|
||||||
|
def hsv_to_color_name(h, s, v):
|
||||||
|
"""Convert HSV to color name"""
|
||||||
|
if v < 50:
|
||||||
|
return "black"
|
||||||
|
if s < 40 and v > 200:
|
||||||
|
return "white"
|
||||||
|
if 0 <= h <= 15 or 165 <= h <= 180:
|
||||||
|
return "red"
|
||||||
|
if 15 < h <= 35:
|
||||||
|
return "orange"
|
||||||
|
if 35 < h <= 50:
|
||||||
|
return "yellow"
|
||||||
|
if 50 < h <= 85:
|
||||||
|
return "green"
|
||||||
|
if 85 < h <= 105:
|
||||||
|
return "cyan"
|
||||||
|
if 105 < h <= 140:
|
||||||
|
return "blue"
|
||||||
|
if 140 < h <= 165:
|
||||||
|
return "purple"
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
|
def check_color_match(dominant_colors, target_color):
|
||||||
|
"""Check if dominant colors match target color"""
|
||||||
|
if not dominant_colors:
|
||||||
|
return False, 0.0
|
||||||
|
|
||||||
|
target_lower = target_color.lower()
|
||||||
|
match_count = 0
|
||||||
|
total = len(dominant_colors)
|
||||||
|
|
||||||
|
for color_hsv in dominant_colors:
|
||||||
|
h, s, v = color_hsv[0], color_hsv[1], color_hsv[2]
|
||||||
|
color_name = hsv_to_color_name(h, s, v)
|
||||||
|
if color_name == target_lower:
|
||||||
|
match_count += 1
|
||||||
|
|
||||||
|
ratio = match_count / total if total > 0 else 0.0
|
||||||
|
return ratio > 0.3, ratio # Match if >30% of dominant colors match
|
||||||
|
|
||||||
|
def search_by_color(appearance_path, video_path, target_color, output_path, max_frames=500):
|
||||||
|
"""Search for people wearing target color"""
|
||||||
|
if not os.path.exists(appearance_path):
|
||||||
|
print(json.dumps({"error": f"appearance.json not found: {appearance_path}"}))
|
||||||
|
return
|
||||||
|
|
||||||
|
with open(appearance_path) as f:
|
||||||
|
appearance = json.load(f)
|
||||||
|
|
||||||
|
frames = appearance.get("frames", [])
|
||||||
|
fps = appearance.get("fps", 30)
|
||||||
|
|
||||||
|
cap = cv2.VideoCapture(video_path)
|
||||||
|
if not cap.isOpened():
|
||||||
|
print(json.dumps({"error": f"Cannot open video: {video_path}"}))
|
||||||
|
return
|
||||||
|
|
||||||
|
results = []
|
||||||
|
frame_count = 0
|
||||||
|
|
||||||
|
for frame_data in frames[:max_frames]:
|
||||||
|
frame_num = frame_data.get("frame", 0)
|
||||||
|
persons = frame_data.get("persons", [])
|
||||||
|
timestamp = frame_data.get("timestamp", 0)
|
||||||
|
|
||||||
|
if not persons:
|
||||||
|
frame_count += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
|
||||||
|
ret, frame = cap.read()
|
||||||
|
if not ret:
|
||||||
|
frame_count += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
frame_h, frame_w = frame.shape[:2]
|
||||||
|
|
||||||
|
for person in persons:
|
||||||
|
bbox = person.get("bbox", {})
|
||||||
|
if not bbox:
|
||||||
|
continue
|
||||||
|
|
||||||
|
x, y = bbox.get("x", 0), bbox.get("y", 0)
|
||||||
|
w, h = bbox.get("width", 0), bbox.get("height", 0)
|
||||||
|
|
||||||
|
# Extract upper body region (clothing area)
|
||||||
|
upper_h = int(h * 0.6) # Upper 60% of person
|
||||||
|
roi_x = max(0, int(x))
|
||||||
|
roi_y = max(0, int(y))
|
||||||
|
roi_w = min(w, frame_w - roi_x)
|
||||||
|
roi_h = min(upper_h, frame_h - roi_y)
|
||||||
|
|
||||||
|
if roi_w < 10 or roi_h < 10:
|
||||||
|
continue
|
||||||
|
|
||||||
|
roi = frame[roi_y:roi_y+roi_h, roi_x:roi_x+roi_w]
|
||||||
|
|
||||||
|
# Get dominant colors
|
||||||
|
hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
|
||||||
|
pixels = hsv.reshape(-1, 3).astype(np.float32)
|
||||||
|
|
||||||
|
if len(pixels) < 10:
|
||||||
|
continue
|
||||||
|
|
||||||
|
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
|
||||||
|
_, labels, centers = cv2.kmeans(pixels, 5, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
|
||||||
|
counts = np.bincount(labels.flatten())
|
||||||
|
dominant = centers[np.argsort(-counts)[:5]].tolist()
|
||||||
|
|
||||||
|
match, confidence = check_color_match(dominant, target_color)
|
||||||
|
|
||||||
|
if match:
|
||||||
|
results.append({
|
||||||
|
"frame": frame_num,
|
||||||
|
"timestamp": round(timestamp, 2),
|
||||||
|
"bbox": bbox,
|
||||||
|
"confidence": round(confidence, 3),
|
||||||
|
"dominant_colors": [[round(c, 1) for c in dc] for dc in dominant[:3]]
|
||||||
|
})
|
||||||
|
|
||||||
|
frame_count += 1
|
||||||
|
|
||||||
|
cap.release()
|
||||||
|
|
||||||
|
# Summary
|
||||||
|
color_names = set()
|
||||||
|
for r in results:
|
||||||
|
for dc in r.get("dominant_colors", []):
|
||||||
|
if len(dc) >= 3:
|
||||||
|
color_names.add(hsv_to_color_name(dc[0], dc[1], dc[2]))
|
||||||
|
|
||||||
|
output = {
|
||||||
|
"file_uuid": os.path.basename(appearance_path).split(".")[0],
|
||||||
|
"target_color": target_color,
|
||||||
|
"total_matches": len(results),
|
||||||
|
"matched_frames": list(set(r["frame"] for r in results)),
|
||||||
|
"results": results[:50], # Limit to 50 results
|
||||||
|
"color_names_found": list(color_names)
|
||||||
|
}
|
||||||
|
|
||||||
|
os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
|
||||||
|
with open(output_path, "w") as f:
|
||||||
|
json.dump(output, f, indent=2)
|
||||||
|
|
||||||
|
print(json.dumps({"success": True, **output}))
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
parser = argparse.ArgumentParser(description="Search for people by clothing color")
|
||||||
|
parser.add_argument("--file-uuid", required=True)
|
||||||
|
parser.add_argument("--color", required=True, choices=list(COLOR_RANGES.keys()))
|
||||||
|
parser.add_argument("--video-path", default="")
|
||||||
|
parser.add_argument("--appearance-path", default="")
|
||||||
|
parser.add_argument("--output", default="")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
output_dir = os.environ.get("MOMENTRY_OUTPUT_DIR", "/Users/accusys/momentry/output")
|
||||||
|
appearance_path = args.appearance_path or f"{output_dir}/{args.file_uuid}.appearance.json"
|
||||||
|
video_path = args.video_path
|
||||||
|
output_path = args.output or f"{output_dir}/{args.file_uuid}.color_search_{args.color}.json"
|
||||||
|
|
||||||
|
if not video_path:
|
||||||
|
# Try to find video in common locations
|
||||||
|
for ext in ["mp4", "mov", "avi"]:
|
||||||
|
candidate = f"/Users/accusys/momentry/var/sftpgo/data/demo/{args.file_uuid}.{ext}"
|
||||||
|
if os.path.exists(candidate):
|
||||||
|
video_path = candidate
|
||||||
|
break
|
||||||
|
|
||||||
|
if not video_path:
|
||||||
|
# Search in output directory for video
|
||||||
|
import glob
|
||||||
|
matches = glob.glob(f"/Users/accusys/momentry/var/sftpgo/**/*{args.file_uuid}*", recursive=True)
|
||||||
|
if matches:
|
||||||
|
video_path = matches[0]
|
||||||
|
|
||||||
|
if not video_path:
|
||||||
|
print(json.dumps({"error": "video_path not found, please provide --video-path"}))
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
search_by_color(appearance_path, video_path, args.color, output_path)
|
||||||
@@ -276,6 +276,15 @@ fn make_tools(pool: &sqlx::PgPool) -> Vec<ToolDef> {
|
|||||||
}),
|
}),
|
||||||
vec!["file_uuid", "node_id"],
|
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_file_info" => tools::exec_get_file_info(pool, &args).await,
|
||||||
"get_representative_frame" => tools::exec_get_representative_frame(pool, &args).await,
|
"get_representative_frame" => tools::exec_get_representative_frame(pool, &args).await,
|
||||||
"analyze_frame" => tools::exec_analyze_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)),
|
_ => Err(format!("Unknown tool: {}", name)),
|
||||||
};
|
};
|
||||||
let content = match result {
|
let content = match result {
|
||||||
|
|||||||
@@ -1092,3 +1092,68 @@ pub async fn exec_tkg_node_detail(
|
|||||||
None => Err("Node not found".to_string()),
|
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