check for gpu

app.py

@@ -35,10 +35,14 @@ onnx_file = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.onnx
 # config = {"PERFORMANCE_HINT": "LATENCY"}
 # compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU", config=config)
 
-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)
+# check if GPU is available
+if torch.cuda.is_available():
+    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)
+else:
+    ort_sess = ort.InferenceSession(onnx_file)
 
 
 class SquarePad:
@@ -55,7 +59,7 @@ def sigmoid(x):
     return 1 / (1 + np.exp(-x))
 
 
-def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, miss_threshold=0.8, marks_threshold=0.6, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
+def inference(x, count_only_api, api_key, img_size=224, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, miss_threshold=0.8, marks_threshold=0.6, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
     print(x)
     #api = HfApi(token=os.environ['DATASET_SECRET'])
     #out_file = str(uuid.uuid1())
@@ -89,9 +93,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, 6))
+    event_type_logits = np.zeros((len(all_frames) + seq_len + stride_length, 7))
     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, 6))
+    event_type_logit_overlaps = np.zeros((len(all_frames) + seq_len + stride_length, 7))
     for _ in range(seq_len + stride_length):  # pad full sequence
         all_frames.append(all_frames[-1])
     batch_list = []
@@ -293,7 +297,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'], 
+    bar = px.bar(x=['single rope speed', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders', 'other'], 
                  y=event_type_probs,
                  template="plotly_dark",
                  title="Event Type Distribution",