MedCat
/

table_detection_3

Model card Files Files and versions Community

nguyenp99 commited on May 3

Commit

22464df

•

1 Parent(s): d749e18

Update engine.py

Files changed (1) hide show

engine.py +1 -1

engine.py CHANGED Viewed

@@ -45,7 +45,7 @@ class Engine(f.Pipe):
         assert len(predictor.results) == 1, 'Only support batch size 1'
         preds = predictor.results[0]
         remove_idc = remove_potiential_table_fp(self.stamp_detector, preds.orig_img, preds.boxes.xyxy, self.cfg.STAMP_REMOVING_IOU_THRESHOLD)
         stamp_removed_preds = remove_box_by_idc(preds.boxes, remove_idc)
         predictor.results[0].boxes = stamp_removed_preds

         assert len(predictor.results) == 1, 'Only support batch size 1'
         preds = predictor.results[0]
+        self.stamp_detector.model.to("cuda:1")
         remove_idc = remove_potiential_table_fp(self.stamp_detector, preds.orig_img, preds.boxes.xyxy, self.cfg.STAMP_REMOVING_IOU_THRESHOLD)
         stamp_removed_preds = remove_box_by_idc(preds.boxes, remove_idc)
         predictor.results[0].boxes = stamp_removed_preds