Update eval.py

eval.py

@@ -6,16 +6,15 @@ import argparse
 import numpy as np
 from tqdm import tqdm
 from skimage.metrics import structural_similarity,peak_signal_noise_ratio
-
 import torch
-
 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 بشكل إجباري لضمان التشغيل الموثوق ***
+DEVICE = torch.device('cpu') 
 
 def dewarp_prompt(img):
     mask = net1_net2_infer_single_im(img,'data/MBD/checkpoint/mbd.pkl')
@@ -26,7 +25,7 @@ def dewarp_prompt(img):
 
 def deshadow_prompt(img):
     h,w = img.shape[:2]
-    # img = cv2.resize(img,(128,128))
+    # img = cv2.resize(img,(1024,1024))
     img = cv2.resize(img,(1024,1024))
     rgb_planes = cv2.split(img)
     result_planes = []
@@ -53,10 +52,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)
@@ -85,11 +84,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)
@@ -98,36 +96,31 @@ 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()
+    # *** تم التعديل: استخدام .float() لـ CPU ***
+    model = model.float() 
     with torch.no_grad():
         pred = model(in_im)
         pred = pred[0][:2].permute(1,2,0).cpu().numpy()
         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):
@@ -137,26 +130,23 @@ 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)
+    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:
@@ -165,9 +155,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
@@ -176,26 +164,23 @@ 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)
+    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:
@@ -204,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)
@@ -216,34 +199,33 @@ 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()
+    in_im = in_im.float()
     with torch.no_grad():
         pred = model(in_im,'binarization')
         pred = pred[:,:2,:,:]
@@ -252,42 +234,36 @@ 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('--dataset', nargs='?', type=str, default='./distorted/',help='Path of input document image')
     args = parser.parse_args()
-    assert args.dataset in all_datasets.keys(), 'Unregisted dataset, dataset must be one of '+', '.join(all_datasets)
+    # يتم تعريف all_datasets لاحقًا في __main__
+    # سنحذف assert مؤقتًا أو نعتمد على تعريفها لاحقًا
+    # assert args.dataset in all_datasets.keys(), 'Unregisted dataset, dataset must be one of '+', '.join(all_datasets)
     return 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       
-    )
-
-    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'])    
+        LayerNorm_type = 'WithBias',
+           dual_pixel_task = True
+           )
+    
+    # تحميل النموذج وتعيينه لـ CPU بشكل إجباري
+    state = convert_state_dict(torch.load(args.model_path, map_location='cpu')['model_state'])
     model.load_state_dict(state)
-
     model.eval()
     model = model.to(DEVICE)
     return model
@@ -310,27 +286,27 @@ def inference_one_im(model,im_path,task):
         cv2.imwrite('./temp.jpg',restorted)
         prompt1,prompt2,prompt3,restorted = appearance(model,'./temp.jpg')
         os.remove('./temp.jpg')
-
     return prompt1,prompt2,prompt3,restorted
 
-
-
 if __name__ == '__main__':
     all_datasets = {'dir300':'dewarping','kligler':'deshadowing','jung':'deshadowing','osr':'deshadowing','docunet_docaligner':'appearance','realdae':'appearance','tdd':'deblurring','dibco18':'binarization'}
-
-    ## model init
-    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # تم تعيين DEVICE بالفعل لـ 'cpu' في بداية الملف. نستخدمه هنا.
     args = get_args()
+    
+    # التأكد من أن مجموعة البيانات المدخلة موجودة
+    assert args.dataset in all_datasets.keys(), 'Unregisted dataset, dataset must be one of '+', '.join(all_datasets)
+    
     model = model_init(args)
-
-    ## inference 
+    
+    ## inference     
     print('Predicting')
     task = all_datasets[args.dataset]
     im_paths = glob.glob(os.path.join('./data/eval/',args.dataset,'*_in.*'))
     for im_path in tqdm(im_paths):
         _,_,_,restorted = inference_one_im(model,im_path,task)
         cv2.imwrite(im_path.replace('_in','_docres'),restorted)
-    
+        
     ## obtain metric
     print('Metric calculating')
     if task == 'dewarping':
@@ -341,22 +317,34 @@ if __name__ == '__main__':
         for im_path in tqdm(im_paths):
             pred = cv2.imread(im_path.replace('_in','_docres'))
             gt = cv2.imread(im_path.replace('_in','_gt'))
+            # لضمان التوافق في الأشكال قبل حساب المقاييس
+            if pred.shape != gt.shape:
+                gt = cv2.resize(gt, (pred.shape[1], pred.shape[0]))
+            
             ssim.append(structural_similarity(pred,gt,multichannel=True))
             psnr.append(peak_signal_noise_ratio(pred, gt))
         print(args.dataset)
         print('ssim:',np.mean(ssim))
         print('psnr:',np.mean(psnr))
+        
     elif task=='binarization':
         fmeasures, pfmeasures,psnrs = [],[],[]
         for im_path in tqdm(im_paths):
             pred = cv2.imread(im_path.replace('_in','_docres'))
             gt = cv2.imread(im_path.replace('_in','_gt'))
+            
+            # لضمان التوافق في الأشكال قبل حساب المقاييس
+            if pred.shape != gt.shape:
+                gt = cv2.resize(gt, (pred.shape[1], pred.shape[0]))
+                
             pred = cv2.cvtColor(pred,cv2.COLOR_BGR2GRAY)
             gt = cv2.cvtColor(gt,cv2.COLOR_BGR2GRAY)
+            
             pred[pred>155]=255
             pred[pred<=155]=0
             gt[gt>155]=255
             gt[gt<=155]=0
+            
             fmeasure, pfmeasure,psnr,_,_,_ = utils.bin_metric(pred,gt)
             fmeasures.append(fmeasure)
             pfmeasures.append(pfmeasure)
@@ -365,5 +353,3 @@ if __name__ == '__main__':
         print('fmeasure:',np.mean(fmeasures))
         print('pfmeasure:',np.mean(pfmeasures))
         print('psnr:',np.mean(psnrs))
-    
-