local_recognizer.py

import os
import shutil
import argparse
from ultralytics import YOLO

# --- Configuration ---
CONFIDENCE_THRESHOLD = 0.5
KEYPOINT_HIP_L = 11
KEYPOINT_HIP_R = 12

def is_normal_photo_ai(model, image_path):
    """Same logic as the cloud version, adapted for local files."""
    try:
        results = model.predict(image_path, verbose=False, conf=0.5)
        if not results: return True
        
        result = results[0]
        num_people = len(result.boxes)
        
        # Keep if it's a group or if AI is unsure
        if num_people == 0: return True 
        if num_people > 1: return True 

        # Exactly one person: check for hips
        keypoints = result.keypoints.data[0]
        hip_l_conf = keypoints[KEYPOINT_HIP_L][2].item()
        hip_r_conf = keypoints[KEYPOINT_HIP_R][2].item()
        
        if hip_l_conf > CONFIDENCE_THRESHOLD or hip_r_conf > CONFIDENCE_THRESHOLD:
            return True 
            
        return False 
    except Exception as e:
        print(f"Error analyzing {image_path}: {e}")
        return True 

def process_local_folders(input_path, output_path):
    print("Loading AI Model...")
    # This will download the weights automatically if missing
    model = YOLO('yolov8n-pose.pt')
    
    # Supported image extensions
    valid_extensions = ('.jpg', '.jpeg', '.png', '.webp')
    
    stats = {"processed": 0, "kept": 0}

    # Walk through the input directory
    for root, dirs, files in os.walk(input_path):
        for file in files:
            if file.lower().endswith(valid_extensions):
                input_file_path = os.path.join(root, file)
                
                # Determine relative path to maintain structure
                rel_path = os.path.relpath(root, input_path)
                dest_dir = os.path.join(output_path, rel_path)
                
                if not os.path.exists(dest_dir):
                    os.makedirs(dest_dir)
                
                dest_file_path = os.path.join(dest_dir, file)
                
                stats["processed"] += 1
                print(f"[{stats['processed']}] Analyzing: {file}...", end="\r")
                
                if is_normal_photo_ai(model, input_file_path):
                    shutil.copy2(input_file_path, dest_file_path)
                    stats["kept"] += 1

    print(f"\n\nDone!")
    print(f"Total Processed: {stats['processed']}")
    print(f"Total Kept:      {stats['kept']}")
    print(f"Output saved to: {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Full Body Recognizer (Local Version)")
    parser.add_argument("--input", required=True, help="Path to the source folder")
    parser.add_argument("--output", required=True, help="Path to the destination folder")
    
    args = parser.parse_args()
    
    # Ensure directories exist
    if not os.path.exists(args.input):
        print(f"Error: Input folder '{args.input}' does not exist.")
    else:
        # Create output if it doesn't exist
        if not os.path.exists(args.output):
            os.makedirs(args.output)
        process_local_folders(args.input, args.output)