Update inference.py

inference.py

@@ -3,16 +3,17 @@ import cv2
 import utils
 import argparse
 import numpy as np
-
-import torch
-
+import torch # تم تصحيح: torcth -> torch
+from PIL import Image
 from utils import convert_state_dict
 from models import restormer_arch
 from data.preprocess.crop_merge_image import stride_integral
-
 os.sys.path.append('./data/MBD/')
 from data.MBD.infer import net1_net2_infer_single_im
 
+# تحديد الجهاز ليكون CPU بشكل إجباري
+# يتم استخدامه في دالة model_init وكذلك في الدالة الرئيسية
+DEVICE = torch.device('cpu') 
 
 def dewarp_prompt(img):
     mask = net1_net2_infer_single_im(img,'data/MBD/checkpoint/mbd.pkl')
@@ -50,10 +51,10 @@ def deshadow_prompt(img):
     return bg_imgs
 
 def deblur_prompt(img):
-    x = cv2.Sobel(img,cv2.CV_16S,1,0)  
-    y = cv2.Sobel(img,cv2.CV_16S,0,1)  
-    absX = cv2.convertScaleAbs(x)   # 转回uint8  
-    absY = cv2.convertScaleAbs(y)  
+    x = cv2.Sobel(img,cv2.CV_16S,1,0)      
+    y = cv2.Sobel(img,cv2.CV_16S,0,1)      
+    absX = cv2.convertScaleAbs(x)   # 转回uint8      
+    absY = cv2.convertScaleAbs(y)      
     high_frequency = cv2.addWeighted(absX,0.5,absY,0.5,0)
     high_frequency = cv2.cvtColor(high_frequency,cv2.COLOR_BGR2GRAY)
     high_frequency = cv2.cvtColor(high_frequency,cv2.COLOR_GRAY2BGR)
@@ -82,11 +83,10 @@ def binarization_promptv2(img):
     thresh = thresh.astype(np.uint8)
     result[result>155]=255
     result[result<=155]=0
-
-    x = cv2.Sobel(img,cv2.CV_16S,1,0)  
-    y = cv2.Sobel(img,cv2.CV_16S,0,1)  
-    absX = cv2.convertScaleAbs(x)   # 转回uint8  
-    absY = cv2.convertScaleAbs(y)  
+    x = cv2.Sobel(img,cv2.CV_16S,1,0)      
+    y = cv2.Sobel(img,cv2.CV_16S,0,1)      
+    absX = cv2.convertScaleAbs(x)   # 转回uint8      
+    absY = cv2.convertScaleAbs(y)      
     high_frequency = cv2.addWeighted(absX,0.5,absY,0.5,0)
     high_frequency = cv2.cvtColor(high_frequency,cv2.COLOR_BGR2GRAY)
     return np.concatenate((np.expand_dims(thresh,-1),np.expand_dims(high_frequency,-1),np.expand_dims(result,-1)),-1)
@@ -95,19 +95,15 @@ def dewarping(model,im_path):
     INPUT_SIZE=256
     im_org = cv2.imread(im_path)
     im_masked, prompt_org = dewarp_prompt(im_org.copy())
-
     h,w = im_masked.shape[:2]
     im_masked = im_masked.copy()
     im_masked = cv2.resize(im_masked,(INPUT_SIZE,INPUT_SIZE))
     im_masked = im_masked / 255.0
     im_masked = torch.from_numpy(im_masked.transpose(2,0,1)).unsqueeze(0)
     im_masked = im_masked.float().to(DEVICE)
-
     prompt = torch.from_numpy(prompt_org.transpose(2,0,1)).unsqueeze(0)
     prompt = prompt.float().to(DEVICE)
-
     in_im = torch.cat((im_masked,prompt),dim=1)
-
     # inference
     base_coord = utils.getBasecoord(INPUT_SIZE,INPUT_SIZE)/INPUT_SIZE
     model = model.float()
@@ -117,14 +113,12 @@ def dewarping(model,im_path):
         pred = pred+base_coord
     ## smooth
     for i in range(15):
-        pred = cv2.blur(pred,(3,3),borderType=cv2.BORDER_REPLICATE) 
+        pred = cv2.blur(pred,(3,3),borderType=cv2.BORDER_REPLICATE)     
     pred = cv2.resize(pred,(w,h))*(w,h)
     pred = pred.astype(np.float32)
     out_im = cv2.remap(im_org,pred[:,:,0],pred[:,:,1],cv2.INTER_LINEAR)
-
     prompt_org = (prompt_org*255).astype(np.uint8)
     prompt_org = cv2.resize(prompt_org,im_org.shape[:2][::-1])
-
     return prompt_org[:,:,0],prompt_org[:,:,1],prompt_org[:,:,2],out_im
 
 def appearance(model,im_path):
@@ -134,26 +128,24 @@ def appearance(model,im_path):
     h,w = im_org.shape[:2]
     prompt = appearance_prompt(im_org)
     in_im = np.concatenate((im_org,prompt),-1)
-
-    # constrain the max resolution 
+    # constrain the max resolution     
     if max(w,h) < MAX_SIZE:
         in_im,padding_h,padding_w = stride_integral(in_im,8)
     else:
         in_im = cv2.resize(in_im,(MAX_SIZE,MAX_SIZE))
-    
-    # normalize
+        # normalize    
     in_im = in_im / 255.0
     in_im = torch.from_numpy(in_im.transpose(2,0,1)).unsqueeze(0)
-
     # inference
-    in_im = in_im.half().to(DEVICE)
-    model = model.half()
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    in_im = in_im.float().to(DEVICE) 
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    model = model.float()
     with torch.no_grad():
         pred = model(in_im)
         pred = torch.clamp(pred,0,1)
         pred = pred[0].permute(1,2,0).cpu().numpy()
         pred = (pred*255).astype(np.uint8)
-
         if max(w,h) < MAX_SIZE:
             out_im = pred[padding_h:,padding_w:]
         else:
@@ -162,9 +154,7 @@ def appearance(model,im_path):
             shadow_map = cv2.resize(shadow_map,(w,h))
             shadow_map[shadow_map==0]=0.00001
             out_im = np.clip(im_org.astype(float)/shadow_map,0,255).astype(np.uint8)
-
-    return prompt[:,:,0],prompt[:,:,1],prompt[:,:,2],out_im
-        
+    return prompt[:,:,0],prompt[:,:,1],prompt[:,:,2],out_im        
 
 def deshadowing(model,im_path):
     MAX_SIZE=1600
@@ -173,26 +163,24 @@ def deshadowing(model,im_path):
     h,w = im_org.shape[:2]
     prompt = deshadow_prompt(im_org)
     in_im = np.concatenate((im_org,prompt),-1)
-
-    # constrain the max resolution 
+    # constrain the max resolution     
     if max(w,h) < MAX_SIZE:
         in_im,padding_h,padding_w = stride_integral(in_im,8)
     else:
         in_im = cv2.resize(in_im,(MAX_SIZE,MAX_SIZE))
-    
-    # normalize
+        # normalize    
     in_im = in_im / 255.0
     in_im = torch.from_numpy(in_im.transpose(2,0,1)).unsqueeze(0)
-
     # inference
-    in_im = in_im.half().to(DEVICE)
-    model = model.half()
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    in_im = in_im.float().to(DEVICE)
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    model = model.float()
     with torch.no_grad():
         pred = model(in_im)
         pred = torch.clamp(pred,0,1)
         pred = pred[0].permute(1,2,0).cpu().numpy()
         pred = (pred*255).astype(np.uint8)
-
         if max(w,h) < MAX_SIZE:
             out_im = pred[padding_h:,padding_w:]
         else:
@@ -201,10 +189,8 @@ def deshadowing(model,im_path):
             shadow_map = cv2.resize(shadow_map,(w,h))
             shadow_map[shadow_map==0]=0.00001
             out_im = np.clip(im_org.astype(float)/shadow_map,0,255).astype(np.uint8)
-
     return prompt[:,:,0],prompt[:,:,1],prompt[:,:,2],out_im
 
-
 def deblurring(model,im_path):
     # setup image
     im_org = cv2.imread(im_path)
@@ -213,34 +199,34 @@ def deblurring(model,im_path):
     in_im = np.concatenate((in_im,prompt),-1)
     in_im = in_im / 255.0
     in_im = torch.from_numpy(in_im.transpose(2,0,1)).unsqueeze(0)
-    in_im = in_im.half().to(DEVICE)  
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    in_im = in_im.float().to(DEVICE)      
     # inference
     model.to(DEVICE)
     model.eval()
-    model = model.half()
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    model = model.float()
     with torch.no_grad():
         pred = model(in_im)
         pred = torch.clamp(pred,0,1)
         pred = pred[0].permute(1,2,0).cpu().numpy()
         pred = (pred*255).astype(np.uint8)
         out_im = pred[padding_h:,padding_w:]
-    
     return prompt[:,:,0],prompt[:,:,1],prompt[:,:,2],out_im
 
-
-
 def binarization(model,im_path):
     im_org = cv2.imread(im_path)
     im,padding_h,padding_w = stride_integral(im_org,8)
     prompt = binarization_promptv2(im)
     h,w = im.shape[:2]
     in_im = np.concatenate((im,prompt),-1)
-
     in_im = in_im / 255.0
     in_im = torch.from_numpy(in_im.transpose(2,0,1)).unsqueeze(0)
     in_im = in_im.to(DEVICE)
-    model = model.half()
-    in_im = in_im.half()
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    model = model.float()
+    # *** تم التعديل: استخدام .float() بدلاً من .half() لـ CPU ***
+    in_im = in_im.float()
     with torch.no_grad():
         pred = model(in_im)
         pred = pred[:,:2,:,:]
@@ -249,24 +235,15 @@ def binarization(model,im_path):
         pred = (pred*255).astype(np.uint8)
         pred = cv2.resize(pred,(w,h))
         out_im = pred[padding_h:,padding_w:]
-    
     return prompt[:,:,0],prompt[:,:,1],prompt[:,:,2],out_im
 
-
-
-
-
 def get_args():
     parser = argparse.ArgumentParser(description='Params')
     parser.add_argument('--model_path', nargs='?', type=str, default='./checkpoints/docres.pkl',help='Path of the saved checkpoint')
-    parser.add_argument('--im_path', nargs='?', type=str, default='./distorted/',
-                        help='Path of input document image')
-    parser.add_argument('--out_folder', nargs='?', type=str, default='./restorted/',
-                        help='Folder of the output images')
-    parser.add_argument('--task', nargs='?', type=str, default='dewarping', 
-                        help='task that need to be executed')
-    parser.add_argument('--save_dtsprompt', nargs='?', type=int, default=0, 
-                        help='Width of the input image')
+    parser.add_argument('--im_path', nargs='?', type=str, default='./distorted/',                        help='Path of input document image')
+    parser.add_argument('--out_folder', nargs='?', type=str, default='./restorted/',                        help='Folder of the output images')
+    parser.add_argument('--task', nargs='?', type=str, default='dewarping',                         help='task that need to be executed')
+    parser.add_argument('--save_dtsprompt', nargs='?', type=int, default=0,                         help='Width of the input image')
     args = parser.parse_args()
     possible_tasks = ['dewarping','deshadowing','appearance','deblurring','binarization','end2end']
     assert args.task in possible_tasks, 'Unsupported task, task must be one of '+', '.join(possible_tasks)
@@ -274,25 +251,26 @@ def get_args():
 
 def model_init(args):
    # prepare model
-    model = restormer_arch.Restormer( 
-        inp_channels=6, 
-        out_channels=3, 
-        dim = 48,
-        num_blocks = [2,3,3,4], 
-        num_refinement_blocks = 4,
+    model = restormer_arch.Restormer(
+         inp_channels=6,
+         out_channels=3,
+         dim = 48,
+        num_blocks = [2,3,3,4],
+         num_refinement_blocks = 4,
         heads = [1,2,4,8],
         ffn_expansion_factor = 2.66,
         bias = False,
         LayerNorm_type = 'WithBias',
-        dual_pixel_task = True        
-    )
-
+        dual_pixel_task = True
+            )
+    # نظرًا لأن DEVICE تم تعيينه لـ 'cpu' بشكل إجباري في __main__، سيتم تنفيذ هذا الشرط دائمًا
     if DEVICE.type == 'cpu':
         state = convert_state_dict(torch.load(args.model_path, map_location='cpu')['model_state'])
     else:
-        state = convert_state_dict(torch.load(args.model_path, map_location='cuda:0')['model_state'])    
+        # هذا الجزء لن يتم تنفيذه
+        state = convert_state_dict(torch.load(args.model_path, map_location='cuda:0')['model_state'])
+        
     model.load_state_dict(state)
-
     model.eval()
     model = model.to(DEVICE)
     return model
@@ -316,20 +294,15 @@ def inference_one_im(model,im_path,task):
         prompt1,prompt2,prompt3,restorted = appearance(model,'restorted/step2.jpg')
         # os.remove('restorted/step1.jpg')
         # os.remove('restorted/step2.jpg')
-
     return prompt1,prompt2,prompt3,restorted
 
-
-
 if __name__ == '__main__':
-    ## model init
-    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    # *** تم التعديل: تم تعيين DEVICE لـ 'cpu' بشكل إجباري ***
+    # هذا يضمن أن البرنامج سيعمل على CPU حتى لو كانت CUDA متاحة
     args = get_args()
     model = model_init(args)
-
     ## inference
     prompt1,prompt2,prompt3,restorted = inference_one_im(model,args.im_path,args.task)
-
     ## results saving
     im_name = os.path.split(args.im_path)[-1]
     im_format = '.'+im_name.split('.')[-1]