Update stamp_processing/detector.py

stamp_processing/detector.py

@@ -30,13 +30,13 @@ class StampDetector:
         print(checkpoint)
 
         self.device = device
-        self.model, self.stride = load_yolo_model(checkpoint, device="cuda:1")
+        self.model, self.stride = load_yolo_model(checkpoint, device=device)
 
         self.img_size = 640
         self.conf_thres = conf_thres
         self.iou_thres = iou_thres
 
-        self.process_func_ = partial(process_image, device="cuda:1")
+        self.process_func_ = partial(process_image, device=device)
 
     def __call__(self, image_list: Union[List[npt.NDArray], npt.NDArray]) -> List[npt.NDArray]:
         """Returns a list of bounding boxes [xmin, ymin, xmax, ymax] for each image in image_list