added different device options

arabic/page_htr.py

@@ -228,7 +228,7 @@ def page_htr_one_file(img_file, config_file, model_mode="pretrain", device="cuda
     return page_json
 
 
-def hw_one_file(img_file, config_file, json_obj, model_mode="pretrain", line_key=None):
+def hw_one_file(img_file, config_file, json_obj, model_mode="pretrain", line_key=None, device="cuda"):
     
     HW, idx_to_char = get_hw(config_file)
 
@@ -243,9 +243,9 @@ def hw_one_file(img_file, config_file, json_obj, model_mode="pretrain", line_key
         img = cv2.imread(img_file)
         line_img = warp.get_line_image(values['coord'], img)
         line_text = test_hw.get_predicted_str(HW, None, idx_to_char, flip=True, 
-                                              img=line_img, read_image=False)
-        line_text = make_manual_text_correction(line_text)
-        # New change
+                                              img=line_img, read_image=False, device=device)
+        
+ 
         line_text_logical_order = clean.get_clean_visual_order(line_text)
         json_obj[line]['text'] = line_text_logical_order       
 
@@ -269,13 +269,14 @@ if __name__ == "__main__":
 
     args = parser.parse_args()
     json_obj = {}
+    device = "cuda" if torch.cuda.is_available() else "cpu"    
     if args.line_htr == 1:
         json_obj = json.loads(args.original_json)
         json_obj = hw_one_file(args.img_path, args.config_file, json_obj, 
-                    model_mode="pretrain", line_key=args.line_key)
+                    model_mode="pretrain", line_key=args.line_key, device=device)
     else:
-        json_obj = json.loads(args.original_json)
-        json_obj = page_htr_one_file(args.img_path, args.config_file, device="cpu")
+        json_obj = json.loads(args.original_json)            
+        json_obj = page_htr_one_file(args.img_path, args.config_file, device=device)
         
     print('BEGIN_OUT')
     print(json.dumps(json_obj))

py3/utils/safe_load.py

@@ -11,7 +11,8 @@ def torch_state(path):
                 print('Warning: Model not found')
             else:
                 print('Good: Model is found')
-            state = torch.load(path, weights_only=True, map_location="cpu")
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            state = torch.load(path, weights_only=True, map_location=device)
             return state
         except:
             print("Failed to load",i,path)