normalize_lms.py

#################################

import sys
from sqlalchemy import create_engine, text,func, select, delete, and_, or_
from sqlalchemy.orm import sessionmaker,scoped_session, declarative_base
from sqlalchemy.pool import NullPool
# from my_declarative_base import Images,ImagesBackground, SegmentTable, Site 
from mp_db_io import DataIO
import pickle
import numpy as np
from pick import pick
import threading
import queue
import csv
import os
import cv2
import mediapipe as mp
import shutil
import pandas as pd
import json
from my_declarative_base import Base, Clusters, Encodings, Detections, Images,PhoneBbox, SegmentTable, SegmentBig, Images
#from sqlalchemy.ext.declarative import declarative_base
from mp_sort_pose import SortPose
from tools_clustering import ToolsClustering
import pymongo
from mediapipe.framework.formats import landmark_pb2
from pymediainfo import MediaInfo
import traceback 
import time
import math
import cv2
from sqlalchemy import Column, Integer, ForeignKey
from constants_make_video import *

NOSE_ID=0


Base = declarative_base()
VERBOSE = False
IS_SSD = False
SSD_PATH = None

SKIP_EXISTING = False # Skips images with a normed bbox but that have Images.h - I think only applies to phone bbox
USE_OBJ = False # do objet detections?
SKIP_BODY = False # skip body landmarks. mostly you want to skip when doing obj bbox
                # or are just redoing hands
REPROCESS_HANDS = False # do hands
ACCEPT_EXSISTING_HANDS = True # if true, it will accept existing data and update bool in SQL to tru
REPROCESSED_BODY = True 
REPROCESSED_BODY_DIFF_THRESH = -1.0
IS_SEGMENT_BIG = False # use SegmentBig table. IF False, and IS_SSD is false, it will use Encodings table
TESTING = False # if true, will not do any DB writes
DETECTIONS_ONLY = False

# join on helper to limit scope and use SSD
INNER_JOIN_HELPER = True
# if you do this, you need to use the correct THIS_CLASS_ID 
# that will call an ID_SEGMENT_DICT to produce the table and folder names 


# INNER_JOIN_TABLE = "SegmentHelperObject_67_phone"
# SSD_PATH = "/Volumes/SSD4_Green/segment_images_detected_63_67"
# INNER_JOIN_TABLE = "SegmentHelperObject_45_salad"
# SSD_PATH = "/Volumes/OWC5/segment_images_book_clock_bowl"

# done 45 73 74 81 82 83 86 92 94 96
# next SegmentHelperObject_80_sign SegmentHelper_topic11_business
# SegmentHelperObject_67_phone 
# SegmentHelperObject_41_cup_glass (big)

LIMIT= 40000000
# Initialize the counter
counter = 2000
STATS_PRINT_EVERY = 1000

# Running totals for body-landmark reprocess comparisons.
rows_processed_count = 0
good_nlms_count = 0
bad_nlms_count = 0

THIS_CLASS_ID = 0 # for object bbox normalization
class_token = ID_SEGMENT_DICT.get(THIS_CLASS_ID, None)
ssd = ID_SSD_DICT.get(THIS_CLASS_ID, None)
if THIS_CLASS_ID in ID_FOLDER_DICT: folder_token = ID_FOLDER_DICT[THIS_CLASS_ID]
else: folder_token = class_token 
# SegmentTable_name = 'SegmentOct20'
SegmentTable_name = 'SegmentBig_isface'
if class_token:
    SegmentHelper_name = f'SegmentHelperObject_{class_token}' # TK revisit this for prodution run
    SegmentFolder = f"/Volumes/{ssd}/segment_images_{folder_token}"
    print("SegmentFolder", SegmentFolder)
    # SORT_TYPE = "obj_bbox_fusion"
else: 
    SegmentHelper_name = 'SegmentHelper_oct2025_evens_quarters'
    # SegmentHelper_name = 'SegmentHelper_T11_Oct20_COCO_Custom_evens_quarters'
    # SegmentFolder = "/Volumes/OWC54/segment_images_T4"
    SegmentFolder = None
    # SegmentHelper_name = 'SegmentHelper_T11_Oct20_COCO_Custom'
    # INNER_JOIN_TABLE = "SegmentHelperObject_100_tulip"

    # HAXXXORS THIS_CLASS_ID is commented out below so that it does ALL classes
    # THIS_CLASS_ID = 82 # for object bbox normalization
    # SegmentFolder = "/Volumes/OWC52/segment_images_OWC4"
io = DataIO(IS_SSD, VERBOSE, SSD_PATH)
db = io.db
# io.db["name"] = "stock"
# io.db["name"] = "ministock"
mongo_client = pymongo.MongoClient(io.dbmongo['host'])
mongo_db = mongo_client[io.dbmongo['name']]
mongo_collection = mongo_db[io.dbmongo['collection']]

face_landmarks_collection = mongo_db["encodings"]
bboxnormed_collection = mongo_db["body_landmarks_norm"]
mongo_hand_collection = mongo_db["hand_landmarks"]

# n_phonebbox_collection= mongo_db["bboxnormed_phone"]

# start a timer
start = time.time()

def ensure_unique_index(collection, field_name):
    # List existing indexes on the collection
    indexes = list(collection.list_indexes())

    # Check if the unique index already exists
    for index in indexes:
        if index['key'] == {field_name: 1}:
            if index.get('unique', False):
                print(f"Unique index on '{field_name}' already exists.")
                return
            else:
                # Drop the non-unique index if it exists
                collection.drop_index(index['name'])
                print(f"Non-unique index on '{field_name}' dropped.")

                # Create a unique index on the specified field
                collection.create_index([(field_name, 1)], unique=True)
                print(f"Unique index on '{field_name}' created successfully.")

# Ensure unique index on the image_id field
ensure_unique_index(bboxnormed_collection, 'image_id')

# Create a database engine
if db['unix_socket']:
    # for MM's MAMP config
    engine = create_engine("mysql+pymysql://{user}:{pw}@/{db}?unix_socket={socket}".format(
        user=db['user'], pw=db['pass'], db=db['name'], socket=db['unix_socket']
    ), pool_pre_ping=True, pool_recycle=600, poolclass=NullPool)
else:
    engine = create_engine("mysql+pymysql://{user}:{pw}@{host}/{db}"
                                .format(host=db['host'], db=db['name'], user=db['user'], pw=db['pass']), pool_pre_ping=True, pool_recycle=600, poolclass=NullPool)


image_edge_multiplier = [1.5,1.5,2,1.5] # bigger portrait
image_edge_multiplier_sm = [1.2, 1.2, 1.6, 1.2] # standard portrait
face_height_output = 500
motion = {"side_to_side": False, "forward_smile": True, "laugh": False, "forward_nosmile": False, "static_pose": False, "simple": False}

EXPAND = False
ONE_SHOT = False # take all files, based off the very first sort order.
JUMP_SHOT = False # jump to random file if can't find a run

cfg = {
    'motion': motion,
    'face_height_output': face_height_output,
    'image_edge_multiplier': image_edge_multiplier_sm,
    'EXPAND': EXPAND,
    'ONE_SHOT': ONE_SHOT,
    'JUMP_SHOT': JUMP_SHOT,
    'HSV_CONTROL': None,
    'VERBOSE': VERBOSE,
    'INPAINT': False,
    'SORT_TYPE': 'planar_hands',
    'OBJ_CLS_ID': 0
}
sort = SortPose(config=cfg)
# sort = SortPose(motion, face_height_output, image_edge_multiplier,EXPAND, ONE_SHOT, JUMP_SHOT, HSV_BOUNDS, VERBOSE,INPAINT, SORT_TYPE, OBJ_CLS_ID)

# Create a session
session = scoped_session(sessionmaker(bind=engine))

# Validate helper table exists early so runs fail fast with a clear message.
if SegmentHelper_name:
    helper_exists = session.execute(
        text("SHOW TABLES LIKE :tbl"),
        {"tbl": SegmentHelper_name}
    ).fetchone()
    if not helper_exists:
        raise RuntimeError(
            f"Configured helper table '{SegmentHelper_name}' does not exist in DB '{db['name']}'."
        )


# Number of threads
#num_threads = io.NUMBER_OF_PROCESSES
num_threads = 1

# Batch configuration for efficient bulk updates
BATCH_SIZE = 200  # Commit every N images instead of per-image

# define SegmentHelper 

class SegmentHelper(Base):
    __tablename__ = SegmentHelper_name
    seg_image_id = Column(Integer, primary_key=True, autoincrement=True)
    image_id = Column(Integer, ForeignKey('Images.image_id'))

# class SegmentHelperInnerJoin(Base):
#     __tablename__ = INNER_JOIN_TABLE
#     seg_image_id = Column(Integer, primary_key=True, autoincrement=True)
#     image_id = Column(Integer, ForeignKey('Images.image_id'))


if VERBOSE: print("objects created")



def get_shape(target_image_id):
    ## get the image somehow
    if VERBOSE: print("get_shape target_image_id", target_image_id)
    
    if IS_SEGMENT_BIG:
        select_image_ids_query = (
            select(SegmentBig.site_name_id, SegmentBig.imagename)
            .filter(SegmentBig.image_id == target_image_id)
        )
    elif (not IS_SEGMENT_BIG and not IS_SSD) or (SegmentHelper_name is not None):
        # use images table
        select_image_ids_query = (
            select(Images.site_name_id, Images.imagename)
            .filter(Images.image_id == target_image_id)
        )
    else:
        select_image_ids_query = (
            select(SegmentTable.site_name_id, SegmentTable.imagename)
            .filter(SegmentTable.image_id == target_image_id)
        )
    
    result = session.execute(select_image_ids_query).fetchall()
    print(f"result {target_image_id}", result)
    site_name_id, imagename = result[0]
    site_specific_root_folder = io.folder_list[site_name_id]
    if SegmentHelper_name is not None and IS_SSD and SegmentFolder is not None:
        file = os.path.join(SegmentFolder, os.path.basename(site_specific_root_folder), imagename)
        print("using SegmentHelper_name for path with SegmentFolder", file)
    else:
        file = site_specific_root_folder + "/" + imagename  # os.path.join acting weird, so avoided
        print("using acting weird path", file)
    if VERBOSE: print("get_shape file", file)

    try:
        if io.platform == "darwin":
            media_info = MediaInfo.parse(file, library_file="/opt/homebrew/opt/libmediainfo/lib/libmediainfo.dylib")
            if VERBOSE: print("darwin got media_info")
        else:
            media_info = MediaInfo.parse(file)
    except Exception as e:
        print("Error getting media info, file not found", target_image_id)
        return None, None

    # Try to get image dimensions from MediaInfo
    for track in media_info.tracks:
        # print("track.track_type", track.track_type)
        
        # Check if it's an image track
        if track.track_type == 'Image':
            if VERBOSE: print("track.height, track.width", track.height, track.width)
            if track.height is not None and track.width is not None:
                return track.height, track.width
        
    # If MediaInfo fails, try loading the image with OpenCV as a fallback
    try:
        image = cv2.imread(file)
        if image is not None:
            height, width = image.shape[:2]  # Extract height and width
            if VERBOSE: print("cv2 found dimensions", height, width)
            return height, width
        else:
            print(f"Could not read the image using cv2, file: {file}")
            return None, None
    except Exception as e:
        print(f"Error loading image with cv2 for file: {file}, error: {e}")
        return None, None

    return None, None


def normalize_obj_bbox(obj_bbox,nose_pos,face_height,shape):
    height,width = shape[:2]
    print("obj_bbox type",type(obj_bbox))
    n_obj_bbox=io.unstring_json(obj_bbox)
    n_obj_bbox["right"]=(n_obj_bbox["right"] -nose_pos["x"])/face_height
    n_obj_bbox["left"]=(n_obj_bbox["left"] -nose_pos["x"])/face_height
    n_obj_bbox["top"]=(n_obj_bbox["top"] -nose_pos["y"])/face_height
    n_obj_bbox["bottom"]=(n_obj_bbox["bottom"] -nose_pos["y"])/face_height
    # n_obj_bbox["right"]=(n_obj_bbox["right"]*width -nose_pos["x"])/face_height
    # n_obj_bbox["left"]=(n_obj_bbox["left"]*width -nose_pos["x"])/face_height
    # n_obj_bbox["top"]=(n_obj_bbox["top"]*height -nose_pos["y"])/face_height
    # n_obj_bbox["bottom"]=(n_obj_bbox["bottom"]*height -nose_pos["y"])/face_height
    print(" 📦 n_obj_bbox",n_obj_bbox)

    return n_obj_bbox

def get_landmarks_mongo(image_id):
    if image_id:
        results = mongo_collection.find_one({"image_id": image_id})
        if results:
            try:
                body_landmarks = results['body_landmarks']
            # print("got encodings from mongo, types are: ", type(face_encodings68), type(face_landmarks), type(body_landmarks))
                return unpickle_array(body_landmarks)
            except KeyError:
                print("KeyError, body_landmarks not found for", image_id)
                return None
        else:
            return None
    else:
        return None


def has_existing_hand_landmarks_norm(image_id):
    if image_id is None:
        return False

    # MySQL can surface ids as int/Decimal/str depending on driver settings.
    # Query all likely forms to avoid false negatives on type mismatch.
    candidate_ids = [image_id]
    try:
        candidate_ids.append(int(image_id))
    except Exception:
        pass
    try:
        candidate_ids.append(str(image_id))
    except Exception:
        pass
    try:
        candidate_ids.append(str(int(image_id)))
    except Exception:
        pass
    # Preserve order, remove duplicates.
    candidate_ids = list(dict.fromkeys(candidate_ids))

    # Fetch all likely matches and inspect in Python to avoid edge cases with
    # nested query operators and mixed schema across historical documents.
    cursor = mongo_hand_collection.find({"image_id": {"$in": candidate_ids}})

    docs_checked = 0
    for doc in cursor:
        # print("checking doc", doc)
        docs_checked += 1

        # Legacy or alternate flat payload shape.
        flat_norm = doc.get("hand_landmarks_norm")
        if isinstance(flat_norm, list) and len(flat_norm) > 0:
            return True
        if flat_norm not in (None, {}, "", []):
            return True

        # Canonical nested payload shape.
        for hand_side in ("left_hand", "right_hand"):
            hand_payload = doc.get(hand_side)
            if not isinstance(hand_payload, dict):
                continue
            hand_norm = hand_payload.get("hand_landmarks_norm")
            if isinstance(hand_norm, list) and len(hand_norm) > 0:
                return True
            if hand_norm not in (None, {}, "", []):
                return True

    if (VERBOSE or TESTING) and docs_checked == 0:
        print("has_existing_hand_landmarks_norm no docs for", image_id, "candidate_ids", candidate_ids)

    return False
    
# def get_hand_landmarks_mongo(image_id):
#     if image_id:
#         results = mongo_hand_collection.find_one({"image_id": image_id})
#         if results:
#             hand_landmarks = results['nlms']
#             # print("got encodings from mongo, types are: ", type(face_encodings68), type(face_landmarks), type(body_landmarks))
#             return unpickle_array(hand_landmarks)
#         else:
#             return None
#     else:
#         return None
    
# def insert_n_landmarks(image_id,n_landmarks):
#     nlms_dict = { "image_id": image_id, "nlms": pickle.dumps(n_landmarks) }
#     x = bboxnormed_collection.insert_one(nlms_dict)
#     print("inserted id",x.inserted_id)
#     return


# def insert_n_landmarks(image_id, n_landmarks):
#     start = time.time()
#     nlms_dict = {"image_id": image_id, "nlms": pickle.dumps(n_landmarks)}
#     result = bboxnormed_collection.update_one(
#         {"image_id": image_id},  # filter
#         {"$set": nlms_dict},     # update
#         upsert=True              # insert if not exists
#     )
#     if result.upserted_id:
#         print("Inserted new document with id:", result.upserted_id)
#     else:
#         print("Updated existing document")
#     print("Time to insert:", time.time()-start)
#     return


def insert_n_phone_bbox(image_id, n_phone_bbox, batch_updates):
    # Collect update for batching; avoid per-row ORM roundtrip.
    batch_updates['PhoneBbox'].append({
        'image_id': int(image_id),
        'bbox_26_norm': n_phone_bbox,
    })
    if VERBOSE:
        print("queued phone bbox norm for image_id:", image_id)

def insert_detections_norm_bbox(detection_id, n_bbox, batch_updates):
    # Collect update for batched SQL by detection_id.
    batch_updates['Detections'].append({
        'detection_id': int(detection_id),
        'bbox_norm': n_bbox,
    })
    if VERBOSE:
        print("queued detection bbox_norm for detection_id:", detection_id)


def unpickle_array(pickled_array):
    if pickled_array:
        try:
            # Attempt to unpickle using Protocol 3 in v3.7
            return pickle.loads(pickled_array, encoding='latin1')
        except TypeError:
            # If TypeError occurs, unpickle using specific protocl 3 in v3.11
            # return pickle.loads(pickled_array, encoding='latin1', fix_imports=True)
            try:
                # Set the encoding argument to 'latin1' and protocol argument to 3
                obj = pickle.loads(pickled_array, encoding='latin1', fix_imports=True, errors='strict', protocol=3)
                return obj
            except pickle.UnpicklingError as e:
                print(f"Error loading pickle data: {e}")
                return None
    else:
        return None

def get_face_height_face_lms(target_image_id,bbox, face_landmarks=None):
    # select target image from mongo mongo_collection if not passed through
    if not face_landmarks:
        results = face_landmarks
    else:
        results = face_landmarks_collection.find_one({"image_id": target_image_id})
    if results:
        # set the face height input properties
        if type(bbox)==str: bbox = io.unstring_json(bbox)
        sort.bbox = bbox
        sort.faceLms = pickle.loads(results['face_landmarks'])
        # set the face height
        sort.get_faceheight_data()
        return sort.face_height
    else:
        return None

def get_face_height_bbox(target_image_id):
    if IS_SEGMENT_BIG:
        select_image_ids_query = (
            select(SegmentBig.bbox)
            .filter(SegmentBig.image_id == target_image_id)
        )
    elif not IS_SEGMENT_BIG and not IS_SSD:
        # use Encodings table
        select_image_ids_query = (
            select(Encodings.bbox)
            .filter(Encodings.image_id == target_image_id)
        )
    else:
        select_image_ids_query = (
            select(SegmentTable.bbox)
            .filter(SegmentTable.image_id == target_image_id)
        )
    result = session.execute(select_image_ids_query).fetchall()
    bbox=result[0][0]

    face_height = sort.convert_bbox_to_face_height(bbox)
    return face_height


def get_obj_bbox(target_image_id):
    predicate = text("""
    (
    bbox_norm IS NULL
    OR NOT (
        JSON_EXTRACT(bbox_norm, '$.left') IS NOT NULL
        OR (
        JSON_TYPE(bbox_norm) = 'STRING'
        AND JSON_VALID(CAST(JSON_UNQUOTE(bbox_norm) AS JSON)) = 1
        AND JSON_EXTRACT(CAST(JSON_UNQUOTE(bbox_norm) AS JSON), '$.left') IS NOT NULL
        )
    )
    )
    """)
    select_image_ids_query = (
        select(Detections.detection_id, Detections.bbox, Detections.class_id, Detections.conf)
        .filter(and_(Detections.image_id == target_image_id))
        .filter(predicate)
    )
    
    result = session.execute(select_image_ids_query).fetchall()
    return result

def get_phone_bbox(target_image_id):
    select_image_ids_query = (
        select(PhoneBbox.bbox_26)
        .filter(PhoneBbox.image_id == target_image_id)
    )
    result = session.execute(select_image_ids_query).fetchall()
    phone_bbox=result[0][0]
    if type(phone_bbox)==str:
        phone_bbox=json.loads(phone_bbox)
        if VERBOSE: print("bbox type", type(phone_bbox))  

    return phone_bbox


def compute_landmarks_mad(old_landmarks, new_landmarks):
    """Compute mean absolute difference across overlapping landmark coordinates.
    
    Handles both landmark_pb2.NormalizedLandmarkList (protobuf) and list formats.
    """
    if old_landmarks is None or new_landmarks is None:
        if VERBOSE:
            print(f"  [MAD DEBUG] One input is None: old={old_landmarks is None}, new={new_landmarks is None}")
        return float("inf")

    try:
        # Convert protobuf NormalizedLandmarkList to nested list if needed
        if hasattr(old_landmarks, 'landmark'):  # It's a protobuf object
            old_list = [[lm.x, lm.y, lm.z, lm.visibility] for lm in old_landmarks.landmark]
        else:
            old_list = old_landmarks
        
        if hasattr(new_landmarks, 'landmark'):  # It's a protobuf object
            new_list = [[lm.x, lm.y, lm.z, lm.visibility] for lm in new_landmarks.landmark]
        else:
            new_list = new_landmarks
        
        old_arr = np.asarray(old_list, dtype=float)
        new_arr = np.asarray(new_list, dtype=float)
        
        if VERBOSE:
            print(f"  [MAD DEBUG] old shape={old_arr.shape}, new shape={new_arr.shape}")
        
        if old_arr.ndim < 2 or new_arr.ndim < 2:
            if VERBOSE:
                print(f"  [MAD DEBUG] Invalid ndim: old={old_arr.ndim}, new={new_arr.ndim}")
            return float("inf")

        n_points = min(old_arr.shape[0], new_arr.shape[0])
        n_dims = min(old_arr.shape[1], new_arr.shape[1])
        if n_points == 0 or n_dims == 0:
            if VERBOSE:
                print(f"  [MAD DEBUG] Empty array: n_points={n_points}, n_dims={n_dims}")
            return float("inf")

        delta = np.abs(old_arr[:n_points, :n_dims] - new_arr[:n_points, :n_dims])
        mad = float(np.nanmean(delta))
        if VERBOSE:
            print(f"  [MAD DEBUG] Computed MAD={mad:.6f}")
        return mad
    except Exception as e:
        if VERBOSE:
            print(f"  [MAD DEBUG] Exception during computation: {e}")
        return float("inf")

def calc_nlm(image_id_to_shape, batch_updates):
    global good_nlms_count, bad_nlms_count
    if VERBOSE: print("calc_nlm image_id_to_shape",image_id_to_shape)
    target_image_id = list(image_id_to_shape.keys())[0]
    body_landmarks = None

    # short circuit and just check if norm hands exist already
    if REPROCESS_HANDS and ACCEPT_EXSISTING_HANDS:
        # Query Mongo directly for existing normalized hands to avoid false
        # negatives when multiple docs exist for the same image_id.
        existing_hand_landmarks_norm = has_existing_hand_landmarks_norm(target_image_id)
        print("existing_hand_landmarks_norm", existing_hand_landmarks_norm)
        if existing_hand_landmarks_norm:
            print(f" ☑️ ☑️ ☑️ ACCEPTING EXISTING NORMALIZED HAND LANDMARKS for image_id {target_image_id}, updating SQL to reflect that.")
            # Collect updates for batching
            batch_updates['Encodings'].append({
                'image_id': target_image_id,
                'mongo_hand_landmarks_norm': 1
            })
            if not IS_SEGMENT_BIG:
                batch_updates['SegmentTable'].append({
                    'image_id': target_image_id,
                    'mongo_hand_landmarks_norm': 1
                })
            return



    # TK this needs to be ported to calc body code
    height, width, bbox, face_height, nose_pixel_x, nose_pixel_y = image_id_to_shape[target_image_id]
    bbox = io.unstring_json(bbox) if type(bbox)==str else bbox
    # get the shape of the image if no height in db
    if height and width:
        if VERBOSE: print(target_image_id, "have height,width already",height,width)
    else:
        height,width=get_shape(target_image_id)
        if not height or not width: 
            print(">> IMAGE NOT FOUND,", target_image_id)
            return
        ToolsClustering.store_image_face_data(
            session=session,
            target_image_id=target_image_id,
            image_h=height,
            image_w=width,
            testing=TESTING,
            auto_commit=False,  # OPTIMIZATION: Batch commits instead of individual commits
        )
    sort.h = height
    sort.w = width
    if VERBOSE: print("height,width from DB:",height,width)
    if VERBOSE: print("target_image_id",target_image_id)

    # Fast path: if face geometry is already stored, avoid the expensive
    # face landmark fetch and keep processing with cached nose/face values.
    has_stored_face_geom = (
        face_height is not None
        and nose_pixel_x is not None
        and nose_pixel_y is not None
    )

    if has_stored_face_geom:
        xB, yB = int(nose_pixel_x), int(nose_pixel_y)
        xF, yF = xB, yB
        nose_pixel_pos_face = (xF, yF)
        nose_pixel_pos_body_withviz = {'x': xB, 'y': yB}
        nose_pixel_pos_body = {'x': xB, 'y': yB}
        # Keep body normalization path alive even when skipping full face fetch.
        body_landmarks = get_landmarks_mongo(target_image_id)
        hand_results = None
    else:


        # get all the landmarks
        face_encodings68, face_landmarks, body_landmarks, body_landmarks_normalized, body_landmarks_3D, hand_results = io.get_encodings_mongo(target_image_id)

        if face_landmarks is None:
            print("FACE LANDMARK NOT FOUND 404, bailing for this one ", target_image_id)
            return
        face_height=get_face_height_face_lms(target_image_id,bbox, face_landmarks)
        
        if sort.VERBOSE: print("face_height from lms",face_height)
        nose_pixel_pos_face = sort.get_face_2d_point(1)
        if sort.VERBOSE: print("nose_pixel_pos from face",nose_pixel_pos_face)
        # only do this if the io.get_encodings_mongo didn't return the body landmarks
        if not body_landmarks: body_landmarks=get_landmarks_mongo(target_image_id)
        if VERBOSE:print("body_landmarks",type(body_landmarks), " first few lms:", body_landmarks[:5] if body_landmarks else "None")
        if body_landmarks and type(body_landmarks) == bytes:
            nose_pixel_pos_body_withviz = sort.set_nose_pixel_pos(body_landmarks,[height,width])
            if nose_pixel_pos_body_withviz is None:
                print(f" ❌ ❌ ❌ SKIP image_id {target_image_id}: invalid body landmarks for nose pixel projection. No DB writes.")
                return
            # exit the whole script
            print(f" ✅ ✅ ✅ Converted and saved bodyLms {target_image_id}, with nose pixel", nose_pixel_pos_body_withviz)
            # sys.exit(0)
        elif face_landmarks and type(face_landmarks) == bytes:
            nose_pixel_pos_body_withviz = sort.set_nose_pixel_pos(face_landmarks,[height,width], bbox)
            if nose_pixel_pos_body_withviz is None:
                print(f" ❌ ❌ ❌ SKIP image_id {target_image_id}: invalid face landmarks or missing bbox for nose pixel projection. No DB writes.")
                return
            print(f" ☑️ ☑️ ☑️ Converted and saved faceLms {target_image_id}, with nose pixel", nose_pixel_pos_body_withviz)
        else:
            print(" ❌ BODY LANDMARK NOT FOUND 404, bailing for this one ", target_image_id)
            return
            nose_pixel_pos_body_withviz = nose_pixel_pos_face
        if sort.VERBOSE: print("nose_pixel_pos from body",nose_pixel_pos_body_withviz)
        # hand_results=get_hand_landmarks_mongo(target_image_id)
        # print("hand_results",hand_results)

        #     # for drawing landmarks on test image
        # landmarks_2d = sort.get_landmarks_2d(row['face_landmarks'], list(range(33)), "list")
        # print("landmarks_2d before drawing", landmarks_2d)
        # cropped_image = sort.draw_point(cropped_image, landmarks_2d, index = 0)                    

        # landmarks_2d = sort.get_landmarks_2d(row['face_landmarks'], list(range(420)), "list")
        # cropped_image = sort.draw_point(cropped_image, landmarks_2d, index = 0)                    

        # Extract x and y coordinates
        xF, yF = int(nose_pixel_pos_face[0]), int(nose_pixel_pos_face[1])
        xB, yB = int(nose_pixel_pos_body_withviz['x']), int(nose_pixel_pos_body_withviz['y'])
        

            # gets ride of visiblity
        nose_pixel_pos_body = {'x': xB, 'y': yB}


    ## begin testing stuff

    ## FOR TESTING
    def visualize_landmarks(target_image_id, nose_pixel_pos_face, nose_pixel_pos_body, body_landmarks):
        all_body_landmarks_dict = sort.get_landmarks_2d(body_landmarks, list(range(33)), "dict")
        # query mysql SegmentTable for imagename
        select_image_ids_query = (
            select(SegmentTable.imagename, SegmentTable.site_name_id)
            .filter(SegmentTable.image_id == target_image_id)
        )
        result = session.execute(select_image_ids_query).fetchall()
        imagename=result[0][0]
        site_name_id=result[0][1]

        # open the image with cv2
        image_path = os.path.join(io.ROOT,io.folder_list[site_name_id],imagename)
        print("image_path",image_path)

        # Draw the circle with the converted integer coordinates
        image = cv2.imread(image_path)
        print("image shape",image.shape)
        cv2.circle(image, (int(nose_pixel_pos_face[0]), int(nose_pixel_pos_face[1])), 5, (0, 255, 0), -1)
        cv2.circle(image, (int(nose_pixel_pos_body['x']), int(nose_pixel_pos_body['y'])), 10, (255, 0, 0), -1)

        # iterate through all_body_landmarks_dict and draw the points
        for key, value in all_body_landmarks_dict.items():
            print("value",value)
            x, y = value
            if x < 10 or y < 10:
                cv2.circle(image, (int(x*sort.w), int(y*sort.h)), 5, (0, 0, 255), -1)
            else:
                cv2.circle(image, (int(x), int(y)), 5, (0, 0, 255), -1)

        cv2.imshow("image", image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    # set the nose pixel position in the expected dictionary format

    # visualize_landmarks(target_image_id, nose_pixel_pos_face, nose_pixel_pos_body, body_landmarks)
    
    ## end testing stuff

    # gets ride of visiblity
    nose_pixel_pos_body = {'x': xB, 'y': yB}

    # Calculate Euclidean distance
    distance = math.sqrt((xB - xF)**2 + (yB - yF)**2)

    if distance > 300:
        print(f" >>> TWO NOSES - {target_image_id}. Dist between nose_pixel_pos and nose_pixel_pos_body:", distance)

        # Collect updates for batching
        batch_updates['Encodings'].append({
            'image_id': target_image_id,
            'two_noses': 1
        })
        if not IS_SEGMENT_BIG:
            batch_updates['SegmentTable'].append({
                'image_id': target_image_id,
                'two_noses': 1
            })
        return

    else:
        # store face_height and nose_pixel_pos_body — only if not already set
        if not TESTING:
            ToolsClustering.store_image_face_data(
                session=session,
                target_image_id=target_image_id,
                face_height=face_height,
                nose_pixel_x=xB,
                nose_pixel_y=yB,
                image_h=height,
                image_w=width,
                testing=TESTING,
                auto_commit=False,  # OPTIMIZATION: Batch commits instead of individual commits
            )


    if hand_results:
        if VERBOSE: print("has hand_landmarks")

        hand_landmarks_norm=sort.normalize_hand_landmarks(hand_results,nose_pixel_pos_body,face_height,[height,width])
        # print("hand_landmarks_norm",hand_landmarks_norm)
        sort.update_hand_landmarks_in_mongo(mongo_hand_collection, target_image_id,hand_landmarks_norm)
        # Collect updates for batching
        batch_updates['Encodings'].append({
            'image_id': target_image_id,
            'mongo_hand_landmarks_norm': 1
        })
        if not IS_SEGMENT_BIG:
            batch_updates['SegmentTable'].append({
                'image_id': target_image_id,
                'mongo_hand_landmarks_norm': 1
            })    
        
    if body_landmarks:
        if VERBOSE: print("has body_landmarks")

        ### NORMALIZE LANDMARKS ###
        if VERBOSE: print("nose_pixel_pos",nose_pixel_pos_body)

        if not SKIP_BODY:
            # if type body_landmarks is bytes, unpickle it
            if type(body_landmarks) == bytes:
                body_landmarks = pickle.loads(body_landmarks)
            n_landmarks=sort.normalize_landmarks(body_landmarks,nose_pixel_pos_body,face_height,[height,width])
            if n_landmarks is None:
                if isinstance(body_landmarks, list):
                    print(f"normalize_landmarks returned None for {target_image_id}; retrying after convert_new_mp_to_old_format")
                    try:
                        body_landmarks = sort.convert_new_mp_to_old_format(body_landmarks)
                        n_landmarks = sort.normalize_landmarks(body_landmarks, nose_pixel_pos_body, face_height, [height,width])
                    except Exception as e:
                        print(f"convert_new_mp_to_old_format failed for {target_image_id}: {e}")
                        n_landmarks = None
                if n_landmarks is None:
                    print(f" ⚔️ ⚔️ ⚔️ SKIP image_id {target_image_id}: could not normalize body_landmarks input format. No DB writes.")
                    return

            # if VERBOSE: print("Time to get norm lms:", time.time()-start)
            # start = time.time()

            if VERBOSE: print("about to insert n_landmarks",n_landmarks)

            recalc_flag = 1
            should_write_norm = True
            if REPROCESSED_BODY:
                existing_doc = bboxnormed_collection.find_one({"image_id": target_image_id})
                if existing_doc and existing_doc.get("nlms"):
                    existing_n_landmarks = unpickle_array(existing_doc.get("nlms"))
                    mad = compute_landmarks_mad(existing_n_landmarks, n_landmarks)
                    if VERBOSE:
                        print(f"reprocess compare image_id {target_image_id}: MAD={mad:.6f}")
                    if mad < REPROCESSED_BODY_DIFF_THRESH:
                        print(f" ☑️ ☑️ ☑️ ACCEPTING existing nlms for {target_image_id}: MAD={mad:.6f}")
                        recalc_flag = 0
                        should_write_norm = False
                        with lock:
                            good_nlms_count += 1
                    else:
                        print(f" ♻️ ♻️ ♻️ REJECTING existing nlms & updating with new nlms for {target_image_id}. MAD={mad:.6f}")
                        with lock:
                            bad_nlms_count += 1
                else:
                    recalc_flag = 1
            
            # store the normalized landmarks in mongo
            if not TESTING and should_write_norm:
                if VERBOSE: print("inserting n_landmarks into mongo for image_id", target_image_id)
                sort.insert_n_landmarks(bboxnormed_collection, target_image_id,n_landmarks)

            body_update = {
                'image_id': target_image_id,
                'mongo_body_landmarks_norm': 1
            }
            if REPROCESSED_BODY:
                body_update['mongo_body_landmarks_norm_recalc'] = recalc_flag

            # Collect updates for batching
            batch_updates['Encodings'].append(body_update)
            if not IS_SEGMENT_BIG:
                batch_updates['SegmentTable'].append({
                    'image_id': target_image_id,
                    'mongo_body_landmarks_norm': 1
                })

            if VERBOSE: print(f"insert_n_landmarks done (should_write_norm = {should_write_norm}) going to get phone bbox")
            # if VERBOSE: print("Time to get insert:", time.time()-start)
            # start = time.time()
        
        if USE_OBJ: 
            obj_results = get_obj_bbox(target_image_id)
            # itterate through the obj_results
            print("obj_results",obj_results)
            for detection_id, obj_bbox, class_id, conf in obj_results:
                if obj_bbox and (obj_bbox != 'null' or conf > 0):
                    print("going to normalize obj_bbox",obj_bbox)
                    n_obj_bbox=normalize_obj_bbox(obj_bbox,nose_pixel_pos_body,face_height,[height,width])
                    # temp comment
                    insert_detections_norm_bbox(detection_id, n_obj_bbox, batch_updates)
                else:
                    print("PHONE BBOX NOT FOUND 404", target_image_id)

            # phone_bbox=get_phone_bbox(target_image_id)
            # if phone_bbox:
            #     n_phone_bbox=normalize_obj_bbox(phone_bbox,nose_pixel_pos_body,face_height,[height,width])
            #     insert_detections_norm_bbox(target_image_id,n_phone_bbox)
            # else:
            #     print("PHONE BBOX NOT FOUND 404", target_image_id)
    else:
        print("BODY LANDMARK NOT FOUND 404", target_image_id)

    ## FOR TESTING
    # projected_landmarks = sort.project_normalized_landmarks(n_landmarks, nose_pixel_pos_body, face_height, [height, width])
    # visualize_landmarks(target_image_id, nose_pixel_pos_face, nose_pixel_pos_body, projected_landmarks)



    return

#######MULTI THREADING##################
# Create a lock for thread synchronization
lock = threading.Lock()
threads_completed = threading.Event()

# Create a queue for distributing work among threads
work_queue = queue.Queue()
        
function=calc_nlm


# old way, with phone bbox
# if USE_OBJ == 26:
#     distinct_image_ids_query = select(Images.image_id.distinct(), Images.h, Images.w, SegmentTable.bbox).\
#         outerjoin(SegmentTable,Images.image_id == SegmentTable.image_id).\
#         outerjoin(PhoneBbox,PhoneBbox.image_id == SegmentTable.image_id).\
#         filter(SegmentTable.bbox != None).\
#         filter(SegmentTable.two_noses.is_(None)).\
#         filter(SegmentTable.mongo_body_landmarks == 1).\
#         filter(PhoneBbox.bbox_26 != None).\
#         filter(PhoneBbox.bbox_26_norm == None).\
#         filter(PhoneBbox.conf_26 != -1).\
#         limit(LIMIT)

# new way, with detections
if USE_OBJ:


    print("doing OBJ using Detections")
    predicate_text = """
        (
            bbox_norm IS NULL
            OR JSON_EXTRACT(bbox_norm, '$.left') IS NULL
        )
        """

    # distinct_image_ids_query = select(Detections.detection_id, Detections.bbox, Detections.class_id, Detections.conf).filter(text(predicate_text))

    distinct_image_ids_query = select(
        Images.image_id.distinct(), 
        Images.h, 
        Images.w, 
        Encodings.bbox,
        Encodings.face_height,
        Encodings.nose_pixel_x,
        Encodings.nose_pixel_y
    ).select_from(Detections).\
    join(Encodings, Encodings.image_id == Detections.image_id).\
    join(Images, Images.image_id == Detections.image_id).\
    filter(Encodings.bbox != None).\
    filter(Encodings.two_noses.is_(None)).\
    filter(Encodings.mongo_body_landmarks == 1).\
    filter(Detections.bbox != None).\
    filter(text(predicate_text)).\
    filter(Detections.conf != -1).\
    limit(LIMIT)    

    # join(SegmentHelper, SegmentHelper.image_id == Detections.image_id).\
    if SegmentHelper is not None and INNER_JOIN_HELPER:
        print(f"SUBSELECT_ON_CLASS_ID: limiting OBJ query to {INNER_JOIN_HELPER} using INNER JOIN on SegmentHelper")
        distinct_image_ids_query = distinct_image_ids_query.\
        join(SegmentHelper, SegmentHelper.image_id == Detections.image_id)
    #filter(Detections.class_id == THIS_CLASS_ID).\
    if not REPROCESSED_BODY:
        distinct_image_ids_query = distinct_image_ids_query.filter(Encodings.mongo_body_landmarks_norm.is_(None))
    if not SKIP_BODY:
        distinct_image_ids_query = distinct_image_ids_query.filter(Encodings.mongo_face_landmarks == 1)
    # distinct_image_ids_query = select(Images.image_id.distinct(), Images.h, Images.w, Encodings.bbox).\
    #     outerjoin(SegmentTable,Images.image_id == SegmentTable.image_id).\
    #     outerjoin(Detections,Detections.image_id == SegmentTable.image_id).\
    #     filter(SegmentTable.bbox != None).\
    #     filter(SegmentTable.two_noses.is_(None)).\
    #     filter(SegmentTable.mongo_body_landmarks == 1).\
    #     filter(Detections.bbox != None).\
    #     filter(Detections.bbox_norm == None).\
    #     filter(Detections.conf != -1).\
    #     limit(LIMIT)

elif REPROCESS_HANDS == True and IS_SEGMENT_BIG == True:
    print("doing HANDS using SegmentTable")
    distinct_image_ids_query = select(Images.image_id.distinct(), Images.h, Images.w, SegmentTable.bbox, Encodings.face_height, Encodings.nose_pixel_x, Encodings.nose_pixel_y).\
        outerjoin(SegmentTable,Images.image_id == SegmentTable.image_id).\
        outerjoin(Encodings, Encodings.image_id == Images.image_id).\
        filter(SegmentTable.bbox != None).\
        filter(SegmentTable.two_noses.is_(None)).\
        filter(SegmentTable.mongo_hand_landmarks == 1).\
        filter(SegmentTable.mongo_hand_landmarks_norm.is_(None)).\
        limit(LIMIT)
elif REPROCESS_HANDS == True and IS_SEGMENT_BIG == False:
    print("doing HANDS using Encodings")
    distinct_image_ids_query = select(Images.image_id.distinct(), Images.h, Images.w, Encodings.bbox, Encodings.face_height, Encodings.nose_pixel_x, Encodings.nose_pixel_y).\
        outerjoin(Images, Images.image_id == Encodings.image_id).\
        filter(Encodings.bbox != None).\
        filter(Encodings.two_noses.is_(None)).\
        filter(Encodings.mongo_hand_landmarks == 1).\
        filter(Encodings.mongo_hand_landmarks_norm.is_(None)).\
        limit(LIMIT)
    if INNER_JOIN_HELPER:
        print(f"SUBSELECT_ON_CLASS_ID: limiting HANDS Encodings query to {INNER_JOIN_HELPER}")
        distinct_image_ids_query = distinct_image_ids_query.\
        limit(None).\
        join(SegmentHelper, SegmentHelper.image_id == Encodings.image_id).\
        limit(LIMIT)

elif not SKIP_BODY and IS_SEGMENT_BIG == True:
    print("doing BODY using SegmentBig")
    distinct_image_ids_query = select(Images.image_id.distinct(), Images.h, Images.w, SegmentTable.bbox, Encodings.face_height, Encodings.nose_pixel_x, Encodings.nose_pixel_y).\
        outerjoin(SegmentTable,Images.image_id == SegmentTable.image_id).\
        outerjoin(Encodings, Encodings.image_id == Images.image_id).\
        filter(SegmentTable.bbox != None).\
        filter(SegmentTable.two_noses.is_(None)).\
        filter(SegmentTable.mongo_body_landmarks == 1).\
        limit(LIMIT)
    if not REPROCESSED_BODY:
        distinct_image_ids_query = distinct_image_ids_query.filter(SegmentTable.mongo_body_landmarks_norm.is_(None))
elif not SKIP_BODY and IS_SEGMENT_BIG == False:
    # use Encodings table
    print("doing BODY using Encodings")
    distinct_image_ids_query = select(Images.image_id.distinct(), Images.h, Images.w, Encodings.bbox, Encodings.face_height, Encodings.nose_pixel_x, Encodings.nose_pixel_y).\
        outerjoin(Images, Images.image_id == Encodings.image_id).\
        filter(Encodings.bbox != None).\
        filter(Encodings.is_face == 1).\
        filter(Encodings.two_noses.is_(None)).\
        filter(Encodings.mongo_body_landmarks == 1).\
        limit(LIMIT)
    if not REPROCESSED_BODY:
        distinct_image_ids_query = distinct_image_ids_query.filter(Encodings.mongo_body_landmarks_norm.is_(None))
    else:
        # skip the ones with a mongo_body_landmarks_norm_recalc value
        distinct_image_ids_query = distinct_image_ids_query.filter(Encodings.mongo_body_landmarks_norm_recalc.is_(None))
    
    if DETECTIONS_ONLY: 
        print("DETECTIONS_ONLY: limiting BODY query to images with Detections")
        distinct_image_ids_query = distinct_image_ids_query.\
        limit(None).\
        join(Detections, Detections.image_id == Encodings.image_id).\
        filter(Detections.bbox != None).\
        filter(Detections.conf > 0).\
        limit(LIMIT)
    if INNER_JOIN_HELPER:
        print(f"SUBSELECT_ON_CLASS_ID: limiting BODY Encodings query to {INNER_JOIN_HELPER}")
        distinct_image_ids_query = distinct_image_ids_query.\
        limit(None).\
        join(SegmentHelper, SegmentHelper.image_id == Encodings.image_id).\
        limit(LIMIT)

        # add this when using a second segment inner join
        # join(SegmentHelperInnerJoin, SegmentHelperInnerJoin.image_id == Encodings.image_id).\

else:
    print(f"doing something else that wasn't caught because SKIP_BODY is {SKIP_BODY}, REPROCESS_HANDS is {REPROCESS_HANDS} and IS_SEGMENT_BIG is {IS_SEGMENT_BIG}")


# # TESTING OVERRIDE seghelper is for testing
# distinct_image_ids_query = select(Images.image_id.distinct(), Images.h, Images.w, SegmentBig.bbox).\
#     outerjoin(SegmentBig,Images.image_id == SegmentBig.image_id).\
#     outerjoin(SegmentHelper,Images.image_id == SegmentHelper.image_id).\
#     outerjoin(Encodings, Encodings.image_id == SegmentBig.image_id).\
#     outerjoin(NMLImages, NMLImages.image_id == SegmentBig.image_id).\
#     filter(Encodings.bbox != None).\
#     filter(Encodings.two_noses.is_(None)).\
#     filter(Encodings.mongo_hand_landmarks == 1).\
#     filter(Encodings.mongo_hand_landmarks_norm.is_(None)).\
#     filter(NMLImages.nml_id > 4191363).\
#     limit(LIMIT)

# put this back in at future date if needed
        # filter(Images.h == None).\
        # filter(SegmentTable.image_id >= 9942966).\

if SKIP_EXISTING:
    print("skipping existing normalized bboxes")
    # skips the ones that have obj bbox which have already been done
    normed_image_ids_query = select(SegmentTable.image_id.distinct(), Images.h, Images.w).\
        outerjoin(Images, Images.image_id == SegmentTable.image_id).\
        outerjoin(PhoneBbox,PhoneBbox.image_id == SegmentTable.image_id).\
        filter(
            or_(
                PhoneBbox.bbox_63_norm != None,
                PhoneBbox.bbox_67_norm != None,
                PhoneBbox.bbox_26_norm != None,
                PhoneBbox.bbox_26_norm != None,
                PhoneBbox.bbox_32_norm != None
            )
        ).\
        filter(SegmentTable.image_id >= 9942966).\
        limit(LIMIT)

    distinct_image_ids_query = distinct_image_ids_query.except_(normed_image_ids_query)


if VERBOSE: 
    # Compile the query with literal values for MySQL Workbench
    compiled_query = distinct_image_ids_query.compile(compile_kwargs={"literal_binds": True})
    workbench_query = str(compiled_query)
    print("=" * 100)
    print("WORKBENCH QUERY (ready to paste into MySQL):")
    print("=" * 100)
    print(workbench_query)
    print("=" * 100)

io.print_sqlalchemy_query(engine,distinct_image_ids_query)
results = session.execute(distinct_image_ids_query).fetchall()
print("query executed, results length", len(results))
# make a dictionary of image_id to shape
for result in results:
    if VERBOSE: print("result", result)
    image_id_to_shape = {}
    image_id, height, width, bbox, face_height, nose_pixel_x, nose_pixel_y = result
    image_id_to_shape[image_id] = (height, width, bbox, face_height, nose_pixel_x, nose_pixel_y)

    ### temp single thread for debugging
    # calc_nlm(image_id_to_shape, lock=None, session=session)
    if VERBOSE: print("done with single thread", image_id_to_shape)
    # print(" ")
    # print(" ")
    work_queue.put(image_id_to_shape)        

# distinct_image_ids = [row[0] for row in session.execute(distinct_image_ids_query).fetchall()]
# if VERBOSE: print("query length",len(distinct_image_ids))
# if VERBOSE: print("distinct_image_ids",(distinct_image_ids))
# for counter,target_image_id in enumerate(distinct_image_ids):
#     if counter%1000==0:print("###########"+str(counter)+"images processed ##########")
    # work_queue.put(target_image_id)        

if VERBOSE: print("queue filled")

def commit_batch_updates(batch_updates):
    """Commit all collected updates using batched WHERE IN updates for efficiency.
    Uses UPDATE ... WHERE image_id IN (...) instead of bulk_update_mappings since
    Encodings/SegmentTable PKs differ from image_id."""
    if not TESTING:
        try:
            from collections import defaultdict
            from sqlalchemy import update as sa_update, bindparam

            def _batch_update_by_image_id(model, records):
                """Group records by (column, value) and issue one UPDATE per group."""
                if not records:
                    return
                col_val_to_ids = defaultdict(list)
                for rec in records:
                    img_id = rec['image_id']
                    for col, val in rec.items():
                        if col != 'image_id':
                            col_val_to_ids[(col, val)].append(img_id)
                for (col_name, val), img_ids in col_val_to_ids.items():
                    session.execute(
                        sa_update(model)
                        .where(model.image_id.in_(img_ids))
                        .values({col_name: val})
                        .execution_options(synchronize_session=False)
                    )

            def _batch_update_detections(records):
                """Update Detections using executemany by detection_id."""
                if not records:
                    return
                params = [
                    {
                        'b_detection_id': rec['detection_id'],
                        'b_bbox_norm': json.dumps(rec['bbox_norm']),
                    }
                    for rec in records
                ]
                stmt = text(
                    "UPDATE Detections "
                    "SET bbox_norm = :b_bbox_norm "
                    "WHERE detection_id = :b_detection_id"
                )
                session.execute(stmt, params)

            _batch_update_by_image_id(Encodings, batch_updates['Encodings'])
            _batch_update_by_image_id(SegmentTable, batch_updates['SegmentTable'])
            _batch_update_detections(batch_updates['Detections'])
            _batch_update_by_image_id(PhoneBbox, batch_updates['PhoneBbox'])

            # Single commit for all updates (critical optimization)
            session.commit()
        except Exception as e:
            session.rollback()
            print(f"Error during batch commit: {e}")
            raise

def threaded_fetching():
    """Process work queue with per-thread batch accumulation."""
    global rows_processed_count
    batch_updates = {
        'Encodings': [],
        'SegmentTable': [],
        'Detections': [],
        'PhoneBbox': []
    }
    batch_count = 0
    
    while not work_queue.empty():
        param = work_queue.get()
        # Process image and collect updates
        function(param, batch_updates)

        with lock:
            rows_processed_count += 1
            if rows_processed_count % STATS_PRINT_EVERY == 0:
                print(
                    f"[nlms running total] rows_processed={rows_processed_count} "
                    f"good_nlms={good_nlms_count} bad_nlms={bad_nlms_count}"
                )

        batch_count += 1
        
        # Commit batch at regular intervals to avoid memory bloat
        if batch_count >= BATCH_SIZE:
            commit_batch_updates(batch_updates)
            # Reset batch for next round
            batch_updates = {
                'Encodings': [],
                'SegmentTable': [],
                'Detections': [],
                'PhoneBbox': []
            }
            batch_count = 0
            if VERBOSE:
                with lock:
                    global counter
                    print(f"Batch committed. Approx {counter} images remaining")
        
        work_queue.task_done()
    
    # Commit remaining updates
    if batch_count > 0:
        commit_batch_updates(batch_updates)

def threaded_processing():
    thread_list = []
    for _ in range(num_threads):
        thread = threading.Thread(target=threaded_fetching)
        thread_list.append(thread)
        thread.start()
    # Wait for all threads to complete
    for thread in thread_list:
        thread.join()
    # Set the event to signal that threads are completed
    threads_completed.set()



threaded_processing()
# Commit the changes to the database
threads_completed.wait()

print("done")
print(
    f"[nlms final total] rows_processed={rows_processed_count} "
    f"good_nlms={good_nlms_count} bad_nlms={bad_nlms_count}"
)
# Close the session
session.close()

# Print the time taken
print("Time taken:", time.time()-start)