new model

app.py

@@ -23,17 +23,31 @@ import torchvision.transforms.functional as F
 from huggingface_hub import hf_hub_download
 from huggingface_hub import HfApi
 
+
+
 plt.style.use('dark_background')
 
-hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.bin", repo_type="model", token=os.environ['DATASET_SECRET'])
-model_xml = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.xml", repo_type="model", token=os.environ['DATASET_SECRET'])
-hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.mapping", repo_type="model", token=os.environ['DATASET_SECRET'])
+checkpoint = hf_hub_download(repo_id="dylanplummer/ropenet", filename="ropenet_keypoint_0.pt", repo_type="model", token=os.environ['DATASET_SECRET'])
+model_file = checkpoint = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.py", repo_type="model", token=os.environ['DATASET_SECRET'])
+os.move(model_file, "model.py")
+from model import RepNet
+# model_xml = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.xml", repo_type="model", token=os.environ['DATASET_SECRET'])
+# hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.mapping", repo_type="model", token=os.environ['DATASET_SECRET'])
 #model_xml = "model_ir/model.xml"
 
-ie = Core()
-model_ir = ie.read_model(model=model_xml)
-config = {"PERFORMANCE_HINT": "LATENCY"}
-compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU", config=config)
+# ie = Core()
+# model_ir = ie.read_model(model=model_xml)
+# config = {"PERFORMANCE_HINT": "LATENCY"}
+# compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU", config=config)
+
+img_size = 224
+backbone = 'mobilenetv3'
+embedding_size = 196
+n_layers_lstm = 1
+separate_rope = False
+save_realtime = False
+model = RepNet(64, backbone=backbone, backbone_scale='0', trainable_backbone=False, distill_frame_model=save_realtime, img_size=img_size, embedding_size=embedding_size, separate_rope=separate_rope)
+
 
 
 class SquarePad:
@@ -80,11 +94,6 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
         frame_i += 1
     cap.release()
 
-    # Get output layer
-    output_layer_period_length = compiled_model_ir.output(0)
-    output_layer_periodicity = compiled_model_ir.output(1)
-    output_layer_marks = compiled_model_ir.output(2)
-    output_layer_event_type = compiled_model_ir.output(3)
     length = len(all_frames)
     period_lengths = np.zeros(len(all_frames) + seq_len + stride_length)
     periodicities = np.zeros(len(all_frames) + seq_len + stride_length)
@@ -129,11 +138,11 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
         idx_list.append(i)
         if len(batch_list) == batch_size:
             batch_X = torch.cat(batch_list)
-            result = compiled_model_ir(batch_X)
-            y1pred = result[output_layer_period_length]
-            y2pred = result[output_layer_periodicity]
-            y3pred = result[output_layer_marks]
-            y4pred = result[output_layer_event_type]
+            result = model(batch_X)
+            y1pred = result[0]
+            y2pred = result[1]
+            y3pred = result[2]
+            y4pred = result[3]
             for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
                 periodLength = y1.squeeze()
                 periodicity = y2.squeeze()
@@ -152,11 +161,11 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
             batch_list.append(batch_list[-1])
             idx_list.append(idx_list[-1])
         batch_X = torch.cat(batch_list)
-        result = compiled_model_ir(batch_X)
-        y1pred = result[output_layer_period_length]
-        y2pred = result[output_layer_periodicity]
-        y3pred = result[output_layer_marks]
-        y4pred = result[output_layer_event_type]
+        result = model(batch_X)
+        y1pred = result[0]
+        y2pred = result[1]
+        y3pred = result[2]
+        y4pred = result[3]
         for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
             periodLength = y1.squeeze()
             periodicity = y2.squeeze()
@@ -310,10 +319,7 @@ DESCRIPTION += '\nDemo created by [Dylan Plummer](https://dylan-plummer.github.i
 
 with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
     gr.Markdown(DESCRIPTION)
-    with gr.Column():
-        #with gr.Row():
-        in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', width='50%', min_width=400, interactive=True, container=True)
-        placeholder = gr.Markdown(label="", elem_id='placeholder-text')
+    in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', width=400, height=400, interactive=True, container=True)
             
     with gr.Row():
         run_button = gr.Button(value="Run", elem_id='run-button', scale=1)