Create infer.py

Author: zuru

infer.py

@@ -0,0 +1,54 @@
+import argparse
+import os
+import sys
+import numpy
+import cv2
+
+import torch
+
+import models
+import utils
+import exporters
+
+def parse_arguments(args):
+    usage_text = (
+        "Semi-supervised Spherical Depth Estimation Testing."        
+    )
+    parser = argparse.ArgumentParser(description=usage_text)    
+    parser.add_argument("--input_path", type=str, help="Path to the input spherical panorama image.")
+    parser.add_argument('--weights', type=str, help='Path to the trained weights file.')    
+    parser.add_argument('-g','--gpu', type=str, default='0', help='The ids of the GPU(s) that will be utilized. (e.g. 0 or 0,1, or 0,2). Use -1 for CPU.')        
+    return parser.parse_known_args(args)
+
+if __name__ == "__main__":
+    args, unknown = parse_arguments(sys.argv)
+    gpus = [int(id) for id in args.gpu.split(',') if int(id) >= 0]
+    # device & visualizers
+    device = torch.device("cuda:{}" .format(gpus[0])\
+        if torch.cuda.is_available() and len(gpus) > 0 and gpus[0] >= 0\
+        else "cpu")    
+    # model    
+    model = models.get_model("resnet_coord", {})
+    utils.init.initialize_weights(model, args.weights, pred_bias=None)
+    model = model.to(device)
+    # test data
+    width, height = 512, 256
+    if not os.path.exists(args.input_path):
+        print("Input image path does not exist (%s)." % args.input_path)
+        exit(-1)
+    img = cv2.imread(args.input_path)
+    h, w, _ = img.shape
+    if h != height and w != width:
+        img = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
+    img = img.transpose(2, 0, 1) / 255.0
+    img = torch.from_numpy(img).float().expand(1, -1, -1, -1)
+    model.eval()
+    with torch.no_grad():                    
+        left_rgb = img.to(device)
+        ''' Prediction '''
+        left_depth_pred = torch.abs(model(left_rgb))
+        exporters.image.save_data(os.path.join(
+            os.path.dirname(args.input_path),
+            os.path.splitext(os.path.basename(
+                args.input_path))[0] + "_depth.exr"),
+            left_depth_pred, scale=1.0)