Update app.py

app.py

@@ -10,11 +10,14 @@ import torch
 from torchvision.ops import box_convert
 from torchvision.transforms.functional import to_tensor
 from torchvision.transforms import GaussianBlur
+import time
 
 from Ambrosia import pre_process_image
-import time
 
 
+
+DEVICE = "cuda" # cpu or cuda
+
 # Define a custom transform for Gaussian blur
 def gaussian_blur(x, p=0.5, kernel_size_min=3, kernel_size_max=20, sigma_min=0.1, sigma_max=3):
     if x.ndim == 4:
@@ -75,7 +78,7 @@ def load_image(image_source):
 od_model = load_model(
     model_checkpoint_path="groundingdino_swint_ogc.pth", 
     model_config_path="GroundingDINO_SwinT_OGC.cfg.py", 
-    device="cpu")
+    device=DEVICE)
 print("Object detection model loaded")
 
 def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"):
@@ -119,6 +122,7 @@ def detect_objects(og_image, model=od_model, prompt="bug . insect", device="cpu"
 # load beetle classifier model
 repo_id="ChristopherMarais/beetle-model-mini"
 bc_model = from_pretrained_fastai(repo_id)
+bc_model.to(DEVICE)
 # get class names
 labels = np.append(np.array(bc_model.dls.vocab), "Unknown")
 print("Classification model loaded")
@@ -127,7 +131,7 @@ def predict_beetle(img):
     print("Detecting & classifying beetles...")
     start_time = time.perf_counter() # record how long it processes
     # Split image into smaller images of detected objects
-    image_lst = detect_objects(og_image=img, model=od_model, prompt="bug . insect", device="cpu")
+    image_lst = detect_objects(og_image=img, model=od_model, prompt="bug . insect", device=DEVICE)
     
     # pre_process = pre_process_image(manual_thresh_buffer=0.15, image = img) # use image_dir if directory of image used
     # pre_process.segment(cluster_num=2, 
@@ -143,7 +147,7 @@ def predict_beetle(img):
     output_lst = []
     img_cnt = len(image_lst)
     for i in range(0,img_cnt):
-        prob_ar = np.array(bc_model.predict(image_lst[i])[2])
+        prob_ar = np.array(bc_model.predict(image_lst[i])[2].to(DEVICE).cpu())
         unkown_prob = unkown_prob_calc(probs=prob_ar, wedge_threshold=0.85, wedge_magnitude=5, wedge='dynamic')
         prob_ar = np.append(prob_ar, unkown_prob)
         prob_ar = np.around(prob_ar*100, decimals=1)