diff --git a/src/core/processor/tkg.rs b/src/core/processor/tkg.rs index ed7da8b..52a6147 100644 --- a/src/core/processor/tkg.rs +++ b/src/core/processor/tkg.rs @@ -872,6 +872,19 @@ pub async fn query_auto_representative_frame( let fd_table = t("face_detections"); let nodes_table = t("tkg_nodes"); let edges_table = t("tkg_edges"); + let video_table = t("videos"); + + let half_frame: i64 = match sqlx::query_scalar::<_, i64>(&format!( + "SELECT COALESCE(total_frames / 2, 0) FROM {} WHERE file_uuid = $1", + video_table + )) + .bind(file_uuid) + .fetch_optional(pool) + .await? + { + Some(f) if f > 0 => f, + _ => i64::MAX, + }; let mains = sqlx::query_as::<_, (i32, String, String, i64)>(&format!( "SELECT i.id, i.uuid::text, i.name, COUNT(fd.id)::bigint \ @@ -931,17 +944,17 @@ pub async fn query_auto_representative_frame( .fetch_optional(pool).await?; if let Some((f,)) = tkg_frame { - Some(f) + if f <= half_frame { Some(f) } else { None } } else { sqlx::query_scalar::<_, i64>(&format!( "SELECT MIN(fd_a.frame_number)::bigint \ FROM {} fd_a \ JOIN {} fd_b ON fd_a.frame_number = fd_b.frame_number \ WHERE fd_a.file_uuid = $1 AND fd_a.identity_id = $2 \ - AND fd_b.identity_id = $3", + AND fd_b.identity_id = $3 AND fd_a.frame_number <= $4", fd_table, fd_table )) - .bind(file_uuid).bind(id_a).bind(id_b) + .bind(file_uuid).bind(id_a).bind(id_b).bind(half_frame) .fetch_optional(pool).await? } } @@ -958,11 +971,12 @@ pub async fn query_auto_representative_frame( sqlx::query_scalar::<_, i64>(&format!( "SELECT frame_number::bigint FROM {} \ WHERE file_uuid = $1 AND identity_id = $2 \ + AND frame_number <= $3 \ ORDER BY (width::float8 * height::float8) * confidence::float8 DESC \ LIMIT 1", fd_table )) - .bind(file_uuid).bind(first_id) + .bind(file_uuid).bind(first_id).bind(half_frame) .fetch_optional(pool).await? } else { None @@ -976,11 +990,12 @@ pub async fn query_auto_representative_frame( sqlx::query_scalar::<_, i64>(&format!( "SELECT frame_number::bigint FROM {} \ WHERE file_uuid = $1 AND identity_id IS NOT NULL \ + AND frame_number <= $2 \ ORDER BY (width::float8 * height::float8) * confidence::float8 DESC \ LIMIT 1", fd_table )) - .bind(file_uuid) + .bind(file_uuid).bind(half_frame) .fetch_optional(pool).await? } };