add confidence as alpha

app.py

@@ -5,7 +5,7 @@
 
 from tkinter import W
 import gradio as gr
-
+from matplotlib import cm
 import torch
 import torchvision
 from dlclive import DLCLive, Processor
@@ -33,11 +33,12 @@ FONTS = {'amiko': "font/Amiko-Regular.ttf",
 Megadet_Models = {'md_v5a': "megadet_model/md_v5a.0.0.pt",
                   'md_v5b': "megadet_model/md_v5b.0.0.pt"}
 
-DLC_models = {'full_cat': "model/DLC_Cat_resnet_50_iteration-0_shuffle-0",
-              'full_dog': "model/DLC_Dog_resnet_50_iteration-0_shuffle-0",
-              'primate_face': "model/DLC_FacialLandmarks_resnet_50_iteration-1_shuffle-1",
-              'full_human': "model/DLC_human_dancing_resnet_101_iteration-0_shuffle-1",
-              'full_macaque': 'model/DLC_monkey_resnet_50_iteration-0_shuffle-1'}                  
+DLC_folders = {'full_cat': "model/DLC_Cat/",
+              'full_dog': "model/DLC_Dog/",
+              'primate_face': "model/DLC_FacialLandmarks/",
+              'full_human': "model/DLC_human_dancing/",
+              'full_macaque': 'model/DLC_monkey/'}
+                 
 DLC_models_list = ['full_cat', 'full_dog','primate_face', 'full_human', 'full_macaque']
 #########################################
 #  Draw keypoints on image
@@ -67,27 +68,34 @@ def draw_keypoints_on_image(image,
     
     """
     # get a drawing context
-    draw = ImageDraw.Draw(image)  
+    draw = ImageDraw.Draw(image,"RGBA")  
 
     im_width, im_height = image.size
     keypoints_x = [k[0] for k in keypoints]
     keypoints_y = [k[1] for k in keypoints]
     alpha = [k[2] for k in keypoints]
+    norm = matplotlib.colors.Normalize(vmin=0, vmax=255)
+
+    names_for_color = [i for i in map_label_id_to_str.keys()]
+    colores = np.linspace(0, 255, num=len(names_for_color),dtype= int)
 
     # adjust keypoints coords if required
     if use_normalized_coordinates:
         keypoints_x = tuple([im_width * x for x in keypoints_x])
         keypoints_y = tuple([im_height * y for y in keypoints_y])
     
-    cmap = matplotlib.cm.get_cmap('hsv')
+    #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 = [round(num*255) for num in list(cmap(i*10))[:3]] #check!
-        round_outline = [round(num*255) for num in list(cmap2(alpha[i]))[:3]]
+        round_fill = list(cm.viridis(norm(colores[i]),bytes=True))#[round(num*255) for num in list(cmap(i))[:3]] #check!
+        if np.isnan(alpha[i]) == False : 
+            round_fill[3] = round(alpha[i] *255)
+        #print(round_fill)
+        #round_outline = [round(num*255) for num in list(cmap2(alpha[i]))[:3]]
         draw.ellipse([(keypoint_x - marker_size, keypoint_y - marker_size),
                       (keypoint_x + marker_size, keypoint_y + marker_size)],
-                       fill=tuple(round_fill), outline= tuple(round_outline), width=2) #fill and outline: [0,255]
+                       fill=tuple(round_fill), outline= 'black', width=1) #fill and outline: [0,255]
 
         # add string labels around keypoints
         if flag_show_str_labels:
@@ -95,7 +103,7 @@ def draw_keypoints_on_image(image,
                                      font_size) 
             draw.text((keypoint_x + marker_size, keypoint_y + marker_size),#(0.5*im_width, 0.5*im_height), #-------
                       map_label_id_to_str[i],
-                      ImageColor.getcolor(keypt_color, "RGB"), # rgb
+                      ImageColor.getcolor(keypt_color, "RGB"), # rgb #
                       font=font)
                     
 ############################################
@@ -199,7 +207,7 @@ def predict_dlc(list_np_crops,
         # add kpts of this crop to list 
         list_kpts_per_crop.append(keypts_xyp)
         all_kypts.append(keypts_xyp)
-    #return confidence here
+    
     return list_kpts_per_crop
 
 
@@ -221,8 +229,8 @@ def predict_pipeline(img_input,
     ## Get DLC model and labels as strings   
     # TODO: make a dict as for megadetector
     # pdb.set_trace()
-    path_to_DLCmodel = DownloadModel(dlc_model_input_str, 'model/')
-    pose_cfg_path = 'model/pose_cfg.yaml'      
+    path_to_DLCmodel = DownloadModel(dlc_model_input_str, str(DLC_folders[dlc_model_input_str]) )
+    pose_cfg_path =  str(DLC_folders[dlc_model_input_str]) +'pose_cfg.yaml'      
     #pdb.set_trece()
     # extract map label ids to strings
     # pose_cfg_dict['all_joints'] is a list of one-element lists,
@@ -266,7 +274,8 @@ def predict_pipeline(img_input,
                                 font_size=font_size,
                                 keypt_color=keypt_color,
                                 marker_size=marker_size)
-        return img_input
+                                    
+        return img_input, []
 
     else:
         # Compute kpts for each crop
@@ -297,13 +306,13 @@ def predict_pipeline(img_input,
             img_background.paste(img_crop, box = tuple([int(t) for t in md_results.xyxy[0][ic,:2]]))
 
 
-            #set trh!! FIXME
+            
             bb_per_animal = md_results.xyxy[0].tolist()[ic]
             pred = md_results.xyxy[0].tolist()[ic][4]
             if bbox_likelihood_th < pred:
                 draw_rectangle_text(img_background, bb_per_animal ,font_style=font_style,font_size=font_size, keypt_color=keypt_color)  
 
-            print(pred)
+            
 
         download_file  = save_results(md_results,list_kpts_per_crop,map_label_id_to_str,bbox_likelihood_th)         
 

save_results.py

@@ -38,7 +38,8 @@ def save_results(md_results, dlc_outputs,map_label_id_to_str,thr,output_file = '
         
         ## info dlc
         kypts = []
-        pdb.set_trace()
+        #pdb.set_trace()
+        print(dlc_outputs[i])
         for s in dlc_outputs[i]:
             #print(s)
             aux1 = []
@@ -46,7 +47,7 @@ def save_results(md_results, dlc_outputs,map_label_id_to_str,thr,output_file = '
                 aux1.append(float(j))
 
             kypts.append(aux1)
-        pdb.set_trace()
+        #pdb.set_trace()
         aux['dlc_pred']  = dict(zip(labels,kypts))
         info['bb_' + str(new_index[i]) ]=aux