merge yolo and scsfm

Author: Stippler

.gitignore

@@ -133,3 +133,4 @@ res/*
 *.mp4
 !res/dataset*
 *.pt
+*.tar

adabins/infer.py

@@ -72,7 +72,7 @@ def __init__(self, dataset='nyu', device='cuda:0'):
             self.max_depth = 10
             self.saving_factor = 1000  # used to save in 16 bit
             model = UnetAdaptiveBins.build(n_bins=256, min_val=self.min_depth, max_val=self.max_depth)
-            pretrained_path = "./res/weights/adabins_nyu.pt"
+            pretrained_path = "./data/weights/adabins_nyu.pt"
         elif dataset == 'kitti':
             raise ValueError("dataset can be only 'nyu' {}".format(dataset))
             self.min_depth = 1e-3

camera_calibration.py

@@ -1,12 +1,16 @@
 # %%
-import random
-import torch
-from math import pi, tan
+import scsfm
+import argparse
+import time
 import cv2
 import numpy as np
-from numpy.linalg import inv
-import matplotlib.pyplot as plt
-import pandas as pd
+import torch
+from skimage.transform import resize as imresize
+import random
+from math import pi, tan
+# from numpy.linalg import inv
+# import matplotlib.pyplot as plt
+# import pandas as pd
 
 ###########################
 # VISUALIZATION FUNCTIONS #
@@ -31,10 +35,12 @@ def draw(frame, imgpts):
     # corner = corners[0].ravel()
     imgpts = np.int32(imgpts).reshape(-1, 2)
     # draw ground floor in green
-    frame = cv2.drawContours(frame, [imgpts[:4]], -1, (0, 255, 0), -3)
+    frame = cv2.drawContours(frame, [imgpts[:4]],
+                             -1, (0, 255, 0), -3)
     # draw pillars in blue color
     for i, j in zip(range(4), range(4, 8)):
-        frame = cv2.line(frame, tuple(imgpts[i]), tuple(imgpts[j]), (255), 1)
+        frame = cv2.line(frame, tuple(imgpts[i]),
+                         tuple(imgpts[j]), (255), 1)
     # draw top layer in red color
     frame = cv2.drawContours(frame, [imgpts[4:]], -1, (0, 0, 255), 1)
     return frame
@@ -43,7 +49,7 @@ def draw(frame, imgpts):
 ####################
 # FRAME PARAMETERS #
 ####################
-frame = cv2.imread('data/blob/blob21.jpg')
+frame = cv2.imread('data/data.png')
 height, width, channels = frame.shape
 print(f'width: {width} height: {height} channels: {channels}')
 gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
@@ -112,8 +118,8 @@ def draw(frame, imgpts):
 
 # Filter by Area.
 blobParams.filterByArea = True
-blobParams.minArea = 10     # minArea may be adjusted to suit for your experiment
-blobParams.maxArea = 100   # maxArea may be adjusted to suit for your experiment
+blobParams.minArea = 1
+blobParams.maxArea = 100
 
 # Filter by Circularity
 blobParams.filterByCircularity = True
@@ -303,3 +309,70 @@ def find_board(keypoints):
                 [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
 
 cv2.imwrite('res/blobs.png', blob_frame)
+
+
+device = torch.device(
+    "cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+
+def load_tensor_image(img, resize=(256, 320)):
+
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32)
+
+    if resize:
+        img = imresize(img, resize)
+
+    img = np.transpose(img, (2, 0, 1))
+    tensor_img = ((torch.from_numpy(img).unsqueeze(
+        0) / 255 - 0.45) / 0.225).to(device)
+    print(f'tensor shape {tensor_img.shape}')
+    return tensor_img
+
+
+def prediction_to_visual(output, shape=(360, 640)):
+    pred_disp = output.cpu().numpy()[0, 0]
+    img = 1 / pred_disp
+    img = imresize(img, shape).astype(np.float32)
+    return img
+
+
+@torch.no_grad()
+def predict_depth():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--source', type=str,
+                        default='res/dataset1.mp4', help='source')
+    parser.add_argument('--output', type=str,
+                        default='res/scfm.csv', help='source')
+    opt = parser.parse_args()
+    print(opt)
+
+    ################
+    # Load DispNet #
+    ################
+    disp_net = scsfm.DispResNet(18, False).to(device)
+    weights = torch.load('data/weights/scfm-nyu2-test.pth.tar')
+    disp_net.load_state_dict(weights['state_dict'])
+    disp_net.eval()
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+    # cap = cv2.VideoCapture('data/data2.mp4')
+    # while True:
+    then = time.time()
+    frame_rgb = frame_rgb[:, 30:510]
+    tgt_img = load_tensor_image(frame_rgb.copy())
+    print_duration(then, 'capture and convert')
+    then = time.time()
+    output = disp_net(tgt_img)
+    print_duration(then, 'inference')
+
+    cv2.imshow('frame', frame_rgb)
+    cv2.imshow('depth', prediction_to_visual(output))
+    cv2.waitKey(-1)
+
+
+def print_duration(then, prefix=''):
+    print(prefix, 'took %.2f ms' % ((time.time() - then) * 1000))
+
+
+predict_depth()

cpopservice/depth/__init__.py

@@ -0,0 +1,71 @@
+# %%
+import argparse
+import time
+import cv2
+import numpy as np
+import torch
+from skimage.transform import resize as imresize
+import models.scsfm as models
+
+device = torch.device(
+    "cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+
+def load_tensor_image(img, resize=(256, 320)):
+    img = img[:, 30:510]
+    if resize:
+        resized_img = imresize(img, resize)
+    cv2.imshow('resized', resized_img)
+
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32)
+
+    if resize:
+        img = imresize(img, resize)
+
+    img = np.transpose(img, (2, 0, 1))
+    tensor_img = ((torch.from_numpy(img).unsqueeze(
+        0) / 255 - 0.45) / 0.225).to(device)
+    print(f'tensor shape {tensor_img.shape}')
+    return tensor_img
+
+
+def prediction_to_visual(output, shape=(360, 640)):
+    pred_disp = output.cpu().numpy()[0, 0]
+    img = 1 / pred_disp
+    # img = imresize(img, shape).astype(np.float32)
+    return img
+
+
+@torch.no_grad()
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--source', type=str,
+                        default='res/dataset1.mp4', help='source')
+    parser.add_argument('--output', type=str,
+                        default='res/scfm.csv', help='source')
+
+    opt = parser.parse_args()
+    print(opt)
+
+    disp_net = models.DispResNet(18, False).to(device)
+    weights = torch.load('data/weights/scfm-nyu2-test.pth.tar')
+    disp_net.load_state_dict(weights['state_dict'])
+    disp_net.eval()
+
+    if opt.source == '0':
+        source = 0
+    else:
+        source = opt.source
+
+
+
+
+def print_duration(then, prefix=''):
+    print(prefix, 'took %.2f ms' % ((time.time() - then) * 1000))
+
+
+if __name__ == '__main__':
+    main()
+
+# %%