add url models

app.py

@@ -18,6 +18,7 @@ import math
 # import json
 import os
 import yaml
+from model.models import DownloadModel
 
 import pdb
 
@@ -37,7 +38,7 @@ DLC_models = {'full_cat': "model/DLC_Cat_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_models_list = ['full_cat', 'full_dog','primate_face', 'full_human', 'full_macaque']
 #########################################
 #  Draw keypoints on image
 def draw_keypoints_on_image(image,
@@ -78,15 +79,15 @@ def draw_keypoints_on_image(image,
         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('viridis')
+    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*len(keypoints_x)))[:3]]
+        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]]
         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=3)
+                       fill=tuple(round_fill), outline= tuple(round_outline), width=2) #fill and outline: [0,255]
 
         # add string labels around keypoints
         if flag_show_str_labels:
@@ -160,7 +161,7 @@ def predict_dlc(list_np_crops,
                 kpts_likelihood_th,
                 DLCmodel,
                 dlc_proc):
-
+    print(DLCmodel)
     # run dlc thru list of crops
     dlc_live = DLCLive(DLCmodel, processor=dlc_proc)
     dlc_live.init_inference(list_np_crops[0])
@@ -198,9 +199,9 @@ 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 = DLC_models[dlc_model_input_str]
-    pose_cfg_path = os.path.join(path_to_DLCmodel,'pose_cfg.yaml')        
-                     
+    path_to_DLCmodel = DownloadModel(dlc_model_input_str, '/home/sabrina/MegaDetector_DeepLabCut/model/')
+    pose_cfg_path = '/home/sabrina/MegaDetector_DeepLabCut/model/pose_cfg.yaml'      
+    #pdb.set_trece()
     # extract map label ids to strings
     # pose_cfg_dict['all_joints'] is a list of one-element lists,
     with open(pose_cfg_path, "r") as stream:
@@ -290,7 +291,7 @@ gr_image_input = gr.inputs.Image(type="pil", label="Input Image")
 
 
 # Models
-gr_dlc_model_input = gr.inputs.Dropdown(choices=list(DLC_models.keys()), # choices
+gr_dlc_model_input = gr.inputs.Dropdown(choices=list(DLC_models_list), # choices
                                         default='full_cat', # default option
                                         type='value', # Type of value to be returned by component. "value" returns the string of the choice selected, "index" returns the index of the choice selected.
                                         label='Select DeepLabCut model')                                                                     
@@ -318,7 +319,7 @@ gr_labels_font_style = gr.inputs.Dropdown(choices=['amiko', 'nature', 'painter',
                                         label='Select keypoint label font')
 gr_slider_font_size = gr.inputs.Slider(5,30,1,8,
                                         label='Set font size')
-gr_slider_marker_size = gr.inputs.Slider(1,50,10,10,
+gr_slider_marker_size = gr.inputs.Slider(1,20,1,5,
                                         label='Set marker size')   
 
 # list of inputs

model/models.py

@@ -0,0 +1,60 @@
+import urllib.request
+import tarfile
+from tqdm import tqdm
+import os
+import yaml
+from ruamel.yaml import YAML
+
+def read_plainconfig(configname):
+    if not os.path.exists(configname):
+        raise FileNotFoundError(
+            f"Config {configname} is not found. Please make sure that the file exists."
+        )
+    with open(configname) as file:
+        return YAML().load(file)
+
+def DownloadModel(modelname, target_dir):
+    """
+    Downloads a DeepLabCut Model Zoo Project
+    """
+    
+    def show_progress(count, block_size, total_size):
+        pbar.update(block_size)
+
+    def tarfilenamecutting(tarf):
+        """' auxfun to extract folder path
+        ie. /xyz-trainsetxyshufflez/
+        """
+        for memberid, member in enumerate(tarf.getmembers()):
+            if memberid == 0:
+                parent = str(member.path)
+                l = len(parent) + 1
+            if member.path.startswith(parent):
+                member.path = member.path[l:]
+                yield member
+
+    neturls = read_plainconfig("/home/sabrina/MegaDetector_DeepLabCut/model/pretrained_model_urls.yaml") #FIXME
+    
+    if modelname in neturls.keys():
+        url = neturls[modelname]
+        print(url)
+        response = urllib.request.urlopen(url)
+        print(
+            "Downloading the model from the DeepLabCut server @Harvard -> Go Crimson!!! {}....".format(
+                url
+            )
+        )
+        total_size = int(response.getheader("Content-Length"))
+        pbar = tqdm(unit="B", total=total_size, position=0)
+        filename, _ = urllib.request.urlretrieve(url, reporthook=show_progress)
+        with tarfile.open(filename, mode="r:gz") as tar:
+            tar.extractall(target_dir, members=tarfilenamecutting(tar))
+    else:
+        models = [
+            fn
+            for fn in neturls.keys()
+            if "resnet_" not in fn and "mobilenet_" not in fn
+        ]
+        print("Model does not exist: ", modelname)
+        print("Pick one of the following: ", models)
+    return target_dir

model/pretrained_model_urls.yaml

@@ -0,0 +1,9 @@
+#Model Zoo from where the Charles flow(s)
+full_human: http://deeplabcut.rowland.harvard.edu/models/DLC_human_fullbody_resnet_101.tar.gz
+full_dog: http://deeplabcut.rowland.harvard.edu/models/DLC_Dog_resnet_50_iteration-0_shuffle-0.tar.gz
+full_cat: http://deeplabcut.rowland.harvard.edu/models/DLC_Cat_resnet_50_iteration-0_shuffle-0.tar.gz
+primate_face: http://deeplabcut.rowland.harvard.edu/models/DLC_primate_face_resnet_50_iteration-1_shuffle-1.tar.gz
+mouse_pupil_vclose: http://deeplabcut.rowland.harvard.edu/models/DLC_mouse_pupil_vclose_resnet_50_iteration-0_shuffle-1.tar.gz
+horse_sideview: http://deeplabcut.rowland.harvard.edu/models/DLC_Horses_resnet_50_iteration-1_shuffle-1.tar.gz
+full_macaque: http://deeplabcut.rowland.harvard.edu/models/DLC_macaque_full_resnet50.tar.gz
+full_cheetah: http://deeplabcut.rowland.harvard.edu/models/DLC_full_cheetah_resnet_152.tar.