Update data/MBD/infer.py

data/MBD/infer.py

@@ -5,21 +5,21 @@ import torch.nn.functional as F
 import glob
 import cv2
 from tqdm import tqdm
-
 import time
 import os 
 from model.deep_lab_model.deeplab import *
 from MBD import mask_base_dewarper
 import time
-
 from utils import cvimg2torch,torch2cvimg
 
+# 1. تحديد الجهاز ليكون CPU بشكل إجباري
+device = torch.device('cpu')
+print(f"PyTorch running on device: {device}")
 
 
-def net1_net2_infer(model,img_paths,args):
-
+def net1_net2_infer(model, img_paths, args):
     ### validate on the real datasets
-    seg_model=model
+    seg_model = model
     seg_model.eval()
     for img_path in tqdm(img_paths):
         if os.path.exists(img_path.replace('_origin','_capture')):
@@ -28,16 +28,18 @@ def net1_net2_infer(model,img_paths,args):
         ### segmentation mask predict
         img_org = cv2.imread(img_path)
         h_org,w_org = img_org.shape[:2]
-        img = cv2.resize(img_org,(448, 448))       
+        img = cv2.resize(img_org,(448, 448))               
         img = cv2.GaussianBlur(img,(15,15),0,0)
         img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
         img = cvimg2torch(img)
-
+        
         with torch.no_grad():
-            pred = seg_model(img.cuda())
+            # التعديل رقم 1: نقل المدخلات إلى الجهاز المحدد (CPU)
+            pred = seg_model(img.to(device))
             mask_pred = pred[:,0,:,:].unsqueeze(1)
             mask_pred = F.interpolate(mask_pred,(h_org,w_org))
-            mask_pred = mask_pred.squeeze(0).squeeze(0).cpu().numpy()
+            # نقل الناتج إلى CPU قبل تحويله إلى NumPy
+            mask_pred = mask_pred.squeeze(0).squeeze(0).cpu().numpy() 
             mask_pred = (mask_pred*255).astype(np.uint8)
             kernel = np.ones((3,3))
             mask_pred = cv2.dilate(mask_pred,kernel,iterations=3)
@@ -59,40 +61,46 @@ def net1_net2_infer(model,img_paths,args):
         # cv2.waitKey(0)
         cv2.imwrite(img_path.replace('_origin','_capture'),dewarp)
         cv2.imwrite(img_path.replace('_origin','_mask_new'),mask_pred)
-
         grid0 = cv2.resize(grid[:,:,0],(128,128))
         grid1 = cv2.resize(grid[:,:,1],(128,128))
         grid = np.stack((grid0,grid1),axis=-1)
         np.save(img_path.replace('_origin','_grid1'),grid)
 
-
 def net1_net2_infer_single_im(img,model_path):
     seg_model = DeepLab(num_classes=1,
                     backbone='resnet',
                     output_stride=16,
                     sync_bn=None,
                     freeze_bn=False)
-    seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
-    seg_model.cuda()
-    checkpoint = torch.load(model_path)
+    
+    # التعديل رقم 2: إزالة DataParallel لأنها مصممة للـ GPU
+    # واستبدالها بتحميل النموذج مباشرة
+    # seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
+    
+    # التعديل رقم 3: نقل النموذج إلى الجهاز المحدد (CPU) بدلاً من .cuda()
+    seg_model.to(device)
+    
+    # تحميل النموذج باستخدام map_location للتأكد من تحميله على CPU
+    checkpoint = torch.load(model_path, map_location=device)
     seg_model.load_state_dict(checkpoint['model_state'])
+    
     ### validate on the real datasets
     seg_model.eval()
+    
     ### segmentation mask predict
     img_org = img
     h_org,w_org = img_org.shape[:2]
-    img = cv2.resize(img_org,(448, 448))       
+    img = cv2.resize(img_org,(448, 448))           
     img = cv2.GaussianBlur(img,(15,15),0,0)
     img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
     img = cvimg2torch(img)
-
+    
     with torch.no_grad():
-        # from torchtoolbox.tools import summary
-        # print(summary(seg_model,torch.rand((1, 3, 448, 448)).cuda())) 59.4M 135.6G
-
-        pred = seg_model(img.cuda())
+        # التعديل رقم 4: نقل المدخلات إلى الجهاز المحدد (CPU)
+        pred = seg_model(img.to(device))
         mask_pred = pred[:,0,:,:].unsqueeze(1)
         mask_pred = F.interpolate(mask_pred,(h_org,w_org))
+        # نقل الناتج إلى CPU قبل تحويله إلى NumPy
         mask_pred = mask_pred.squeeze(0).squeeze(0).cpu().numpy()
         mask_pred = (mask_pred*255).astype(np.uint8)
         kernel = np.ones((3,3))
@@ -100,52 +108,32 @@ def net1_net2_infer_single_im(img,model_path):
         mask_pred = cv2.erode(mask_pred,kernel,iterations=3)
         mask_pred[mask_pred>100] = 255
         mask_pred[mask_pred<100] = 0
-        ### tps transform base on the mask
-        # dewarp, grid = mask_base_dewarper(img_org,mask_pred)
-    # try:
-    #     dewarp, grid = mask_base_dewarper(img_org,mask_pred)
-    # except:
-    #     print('fail')
-    #     grid = np.meshgrid(np.arange(w_org),np.arange(h_org))/np.array([w_org,h_org]).reshape(2,1,1)
-    #     grid = torch.from_numpy((grid-0.5)*2).float().unsqueeze(0).permute(0,2,3,1)
-    #     dewarp = torch2cvimg(F.grid_sample(cvimg2torch(img_org),grid))[0]
-    #     grid = grid[0].numpy()
-    # cv2.imshow('in',cv2.resize(img_org,(512,512)))
-    # cv2.imshow('out',cv2.resize(dewarp,(512,512)))
-    # cv2.waitKey(0)
-    # cv2.imwrite(img_path.replace('_origin','_capture'),dewarp)
-    # cv2.imwrite(img_path.replace('_origin','_mask_new'),mask_pred)
-
-    # grid0 = cv2.resize(grid[:,:,0],(128,128))
-    # grid1 = cv2.resize(grid[:,:,1],(128,128))
-    # grid = np.stack((grid0,grid1),axis=-1)
-    # np.save(img_path.replace('_origin','_grid1'),grid)
+    
     return mask_pred
 
-
-
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Hyperparams')
     parser.add_argument('--img_folder', nargs='?', type=str, default='./all_data',help='Data path to load data')
-    parser.add_argument('--img_rows', nargs='?', type=int, default=448, 
-                        help='Height of the input image')
-    parser.add_argument('--img_cols', nargs='?', type=int, default=448, 
-                        help='Width of the input image')
-    parser.add_argument('--seg_model_path', nargs='?', type=str, default='checkpoints/mbd.pkl',
-                        help='Path to previous saved model to restart from')
+    parser.add_argument('--img_rows', nargs='?', type=int, default=448,                         help='Height of the input image')
+    parser.add_argument('--img_cols', nargs='?', type=int, default=448,                         help='Width of the input image')
+    parser.add_argument('--seg_model_path', nargs='?', type=str, default='checkpoints/mbd.pkl',                        help='Path to previous saved model to restart from')
     args = parser.parse_args()
-
+    
     seg_model = DeepLab(num_classes=1,
                     backbone='resnet',
                     output_stride=16,
                     sync_bn=None,
                     freeze_bn=False)
-    seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
-    seg_model.cuda()
-    checkpoint = torch.load(args.seg_model_path)
+    
+    # التعديل رقم 5: إزالة DataParallel
+    # seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
+    
+    # التعديل رقم 6: نقل النموذج إلى الجهاز المحدد (CPU) بدلاً من .cuda()
+    seg_model.to(device)
+    
+    # تحميل النموذج باستخدام map_location للتأكد من تحميله على CPU
+    checkpoint = torch.load(args.seg_model_path, map_location=device)
     seg_model.load_state_dict(checkpoint['model_state'])
-
+    
     im_paths = glob.glob(os.path.join(args.img_folder,'*_origin.*'))
-
     net1_net2_infer(seg_model,im_paths,args)
-