Align DLC outputs with filtered detections

Carry the filtered MegaDetector bounding boxes through cropping, visualization, and JSON export so pose results stay matched to the correct animal detections. This also removes the obsolete save_results module and hardens DLC inference by handling empty crop lists, converting crop channel order for DLCLive, and masking low-confidence keypoints with NaN.

app.py

@@ -70,10 +70,10 @@ def predict_pipeline(img_input,
                                 size=640) #Image.fromarray(results.imgs[0])
 
         ################################################################
-        # Obtain animal crops for bboxes with confidence above th
-        list_crops = crop_animal_detections(img_input,
-                                            md_results,
-                                            bbox_likelihood_th)
+        # Obtain animal crops (and their bboxes) with confidence above th
+        list_crops, list_bboxes = crop_animal_detections(img_input,
+                                                         md_results,
+                                                         bbox_likelihood_th)
 
         ############################################################
 
@@ -135,8 +135,9 @@ def predict_pipeline(img_input,
                                            md_results.ims[0].shape[0]))
         
         # draw keypoints on each crop and paste to background img
-        for ic, (np_crop, kpts_crop) in enumerate(zip(list_crops,
-                                                      list_kpts_per_crop)):
+        for np_crop, kpts_crop, bb_per_animal in zip(list_crops,
+                                                     list_kpts_per_crop,
+                                                     list_bboxes):
 
             img_crop = Image.fromarray(np_crop)
 
@@ -152,20 +153,17 @@ def predict_pipeline(img_input,
                                     marker_size=marker_size)
 
             # Paste crop in original image
-            img_background.paste(img_crop, 
-                                 box = tuple([int(t) for t in md_results.xyxy[0][ic,:2]]))
+            img_background.paste(img_crop,
+                                 box = tuple([int(t) for t in bb_per_animal[:2]]))
 
             # Plot bbox
-            bb_per_animal = md_results.xyxy[0].tolist()[ic]
-            pred = md_results.xyxy[0].tolist()[ic][4]
-            if bbox_likelihood_th < pred:
-                draw_bbox_w_text(img_background, 
-                                    bb_per_animal,
-                                    font_size=font_size)  # TODO: add selectable color for bbox?
-
-            
+            draw_bbox_w_text(img_background,
+                             bb_per_animal,
+                             font_size=font_size)  # TODO: add selectable color for bbox?
+
+
         # Save detection results as json
-        download_file  = save_results_as_json(md_results,list_kpts_per_crop,map_label_id_to_str, bbox_likelihood_th,dlc_model_input_str,mega_model_input)         
+        download_file  = save_results_as_json(md_results,list_kpts_per_crop,list_bboxes,map_label_id_to_str,dlc_model_input_str,mega_model_input)
 
         return img_background, download_file
 

detection_utils.py

@@ -56,8 +56,9 @@ def crop_animal_detections(img_in,
                            likelihood_th):
 
     ## Extract animal crops
-    list_labels_as_str = [i for i in yolo_results.names.values()]  # ['animal', 'person', 'vehicle'] 
+    list_labels_as_str = [i for i in yolo_results.names.values()]  # ['animal', 'person', 'vehicle']
     list_np_animal_crops = []
+    list_animal_bboxes = [] # detection rows [x1,y1,x2,y2,conf,label] matching each crop
 
     # image to crop (scale as input for megadetector)
     img_in = img_in.resize((yolo_results.ims[0].shape[1],
@@ -88,5 +89,6 @@ def crop_animal_detections(img_in,
 
                 # add to list
                 list_np_animal_crops.append(crop_np)
+                list_animal_bboxes.append(det_array[j,:].tolist())
 
-    return list_np_animal_crops
+    return list_np_animal_crops, list_animal_bboxes

dlc_utils.py

@@ -11,22 +11,24 @@ def predict_dlc(list_np_crops,
                 dlc_model_folder,
                 dlc_proc):
     
+    # no animal detected: nothing to run
+    if len(list_np_crops) == 0:
+        return []
+
+    # DLCLive always converts 3-channel frames BGR->RGB, but our crops are
+    # already RGB (PIL), so pass them as BGR for the model to see RGB
+    list_np_crops = [np.ascontiguousarray(crop[..., ::-1]) for crop in list_np_crops]
+
     # run dlc thru list of crops
     dlc_live = DLCLive(dlc_model_folder, processor=dlc_proc)
     dlc_live.init_inference(list_np_crops[0])
 
     list_kpts_per_crop = []
-    all_kypts  = []
-    np_aux = np.empty((1,3)) # can I avoid hardcoding here?
     for crop in list_np_crops:
-        # scale crop here?
         keypts_xyp = dlc_live.get_pose(crop) # third column is llk!
         # set kpts below threhsold to nan
-        
-        #pdb.set_trace()
-        keypts_xyp[keypts_xyp[:,-1] < kpts_likelihood_th,:] = np_aux.fill(np.nan)
-        # add kpts of this crop to list 
+        keypts_xyp[keypts_xyp[:,-1] < kpts_likelihood_th,:] = np.nan
+        # add kpts of this crop to list
         list_kpts_per_crop.append(keypts_xyp)
-        all_kypts.append(keypts_xyp)
-    
+
     return list_kpts_per_crop

save_results.py

@@ -1,56 +0,0 @@
-import json 
-import numpy as np
-import pdb
-
-dict_pred = {0: 'animal', 1: 'person', 2: 'vehicle'}
-
-
-def save_results(md_results, dlc_outputs,map_label_id_to_str,thr,output_file = 'dowload_predictions.json'):
-
-    """
-
-    write json
-
-    """
-    info = {}
-    ## info megaDetector
-    info['file']= md_results.files[0]
-    number_bb = len(md_results.xyxy[0].tolist())
-    info['number_of_bb'] = number_bb
-    number_bb_thr = len(dlc_outputs)
-    labels = [n for n in map_label_id_to_str.values()]
-    #pdb.set_trace()
-    new_index = []
-    for i in range(number_bb):
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[i]
-
-        if confidence > thr:
-            new_index.append(i)
-
-
-    for i in range(number_bb_thr):
-        aux={}
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[new_index[i]]
-        aux['corner_1'] = (corner_x1,corner_y1)
-        aux['corner_2'] = (corner_x2,corner_y2)
-        aux['predict MD'] = md_results.names[0]
-        aux['confidence MD'] = confidence
-        
-        ## info dlc
-        kypts = []
-        for s in dlc_outputs[i]:
-            aux1 = []
-            for j in s:
-                aux1.append(float(j))
-
-            kypts.append(aux1)
-        aux['dlc_pred']  = dict(zip(labels,kypts))
-        info['bb_' + str(new_index[i]) ]=aux
-
-
-    with open(output_file, 'w') as f:
-        json.dump(info, f, indent=1)
-        print('Output file saved at {}'.format(output_file))
-
-    return output_file
-

viz_utils.py

@@ -62,7 +62,6 @@ def draw_keypoints_on_image(image,
         keypoints_y = tuple([im_height * y for y in keypoints_y])
     
     #cmap = matplotlib.cm.get_cmap('hsv')
-    cmap2 = matplotlib.cm.get_cmap('Greys')
     # draw ellipses around keypoints
     for i, (keypoint_x, keypoint_y) in enumerate(zip(keypoints_x, keypoints_y)):
         round_fill = list(cm.viridis(norm(colores[i]),bytes=True))#[round(num*255) for num in list(cmap(i))[:3]] #check!
@@ -113,11 +112,12 @@ def draw_bbox_w_text(img,
     return imgR
 
 ###########################################
-def save_results_as_json(md_results, dlc_outputs, map_dlc_label_id_to_str, thr,model,mega_model_input, path_to_output_file = 'download_predictions.json'):
+def save_results_as_json(md_results, dlc_outputs, animal_bboxes, map_dlc_label_id_to_str, model,mega_model_input, path_to_output_file = 'download_predictions.json'):
 
     """
     Output detections as json file
 
+    animal_bboxes: detection rows [x1,y1,x2,y2,conf,label], one per entry of dlc_outputs
     """
     # initialise dict to save to json
     info = {}
@@ -128,29 +128,20 @@ def save_results_as_json(md_results, dlc_outputs, map_dlc_label_id_to_str, thr,m
     number_bb = len(md_results.xyxy[0].tolist())
     info['number_of_bb'] = number_bb
     # info from DLC
-    number_bb_thr = len(dlc_outputs)
+    info['dlc_model'] = model
     labels = [n for n in map_dlc_label_id_to_str.values()]
-    
-    # create list of bboxes above th
-    new_index = []
-    for i in range(number_bb):
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[i]
-
-        if confidence > thr:
-            new_index.append(i)
 
-    # define aux dict for every bounding box above threshold
-    for i in range(number_bb_thr):
+    # define aux dict for every animal bounding box above threshold
+    for i in range(len(dlc_outputs)):
         aux={}
         # MD output
-        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  md_results.xyxy[0].tolist()[new_index[i]]
+        corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ =  animal_bboxes[i]
         aux['corner_1'] = (corner_x1,corner_y1)
         aux['corner_2'] = (corner_x2,corner_y2)
         aux['predict MD'] = md_results.names[0]
         aux['confidence MD'] = confidence
         
         # DLC output
-        info['dlc_model'] = model
         kypts = []
         for s in dlc_outputs[i]:
             aux1 = []
@@ -159,7 +150,7 @@ def save_results_as_json(md_results, dlc_outputs, map_dlc_label_id_to_str, thr,m
 
             kypts.append(aux1)
         aux['dlc_pred']  = dict(zip(labels,kypts))
-        info['bb_' + str(new_index[i]) ]=aux
+        info['bb_' + str(i) ]=aux
 
     # save dict as json
     with open(path_to_output_file, 'w') as f: