test onnx

app.py

@@ -16,6 +16,7 @@ from passlib.hash import pbkdf2_sha256
 from tqdm import tqdm
 import pandas as pd
 import plotly.express as px
+import onnxruntime as ort
 import torch
 from torchvision import transforms
 import torchvision.transforms.functional as F
@@ -23,14 +24,11 @@ import torchvision.transforms.functional as F
 from huggingface_hub import hf_hub_download
 from huggingface_hub import HfApi
 
-
+from model import RepNet
 
 plt.style.use('dark_background')
 
-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.rename(model_file, "model.py")
-from model import RepNet
+onnx_file = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.onnx", 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'])
 #model_xml = "model_ir/model.xml"
@@ -40,14 +38,10 @@ from model import RepNet
 # 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)
-
+providers = [("CUDAExecutionProvider", {"device_id": torch.cuda.current_device(),
+                                        "user_compute_stream": str(torch.cuda.current_stream().cuda_stream)})]
+sess_options = ort.SessionOptions()
+ort_sess = ort.InferenceSession(onnx_file, sess_options=sess_options, providers=providers)
 
 
 class SquarePad:
@@ -98,9 +92,9 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
     period_lengths = np.zeros(len(all_frames) + seq_len + stride_length)
     periodicities = np.zeros(len(all_frames) + seq_len + stride_length)
     full_marks = np.zeros(len(all_frames) + seq_len + stride_length)
-    event_type_logits = np.zeros((len(all_frames) + seq_len + stride_length, 7))
+    event_type_logits = np.zeros((len(all_frames) + seq_len + stride_length, 6))
     period_length_overlaps = np.zeros(len(all_frames) + seq_len + stride_length)
-    event_type_logit_overlaps = np.zeros((len(all_frames) + seq_len + stride_length, 7))
+    event_type_logit_overlaps = np.zeros((len(all_frames) + seq_len + stride_length, 6))
     for _ in range(seq_len + stride_length):  # pad full sequence
         all_frames.append(all_frames[-1])
     batch_list = []
@@ -138,11 +132,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 = model(batch_X)
-            y1pred = result[0]
-            y2pred = result[1]
-            y3pred = result[2]
-            y4pred = result[3]
+            outputs = ort_sess.run(None, {'frames': batch_X.numpy()})
+            y1pred = outputs[0]
+            y2pred = outputs[1]
+            y3pred = outputs[2]
+            y4pred = outputs[3]
             for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
                 periodLength = y1.squeeze()
                 periodicity = y2.squeeze()
@@ -161,11 +155,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 = model(batch_X)
-        y1pred = result[0]
-        y2pred = result[1]
-        y3pred = result[2]
-        y4pred = result[3]
+        outputs = ort_sess.run(None, {'frames': batch_X.numpy()})
+        y1pred = outputs[0]
+        y2pred = outputs[1]
+        y3pred = outputs[2]
+        y4pred = outputs[3]
         for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
             periodLength = y1.squeeze()
             periodicity = y2.squeeze()
@@ -240,7 +234,7 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
     jumps_per_second[misses] = 0
     frame_type = np.array(['miss' if miss else 'frame' for miss in misses])
     frame_type[full_marks > marks_threshold] = 'jump'
-    per_frame_event_types = np.clip(per_frame_event_types, 0, 7) / 7
+    per_frame_event_types = np.clip(per_frame_event_types, 0, 6) / 6
     df = pd.DataFrame.from_dict({'period length': periodLength, 
                                  'jumping speed': jumping_speed,
                                 'jumps per second': jumps_per_second,
@@ -302,7 +296,7 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
     
     # make a bar plot of the event type distribution
 
-    bar = px.bar(x=['single rope speed', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders', 'other'], 
+    bar = px.bar(x=['single rope speed', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders'], 
                  y=event_type_probs,
                  template="plotly_dark",
                  title="Event Type Distribution",

requirements.txt

@@ -9,4 +9,4 @@ opencv-python-headless==4.7.0.68
 # openvino-dev==2022.3.0
 torch
 torchvision
-timm
+onnxruntime-gpu