diff --git a/.gitignore b/.gitignore index 723ef36..b435414 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,6 @@ -.idea \ No newline at end of file +.DS_Store +.idea + +*venv/ +*.pkl +*.avi diff --git a/requirements.txt b/requirements.txt index c41093f..4f9a42d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ osqp==0.6.1 cvxpy==1.2.1 matplotlib==3.5.3 numpy==1.23.3 -opencv-contrib-python==4.6.0.66 \ No newline at end of file +opencv-contrib-python==4.6.0.66 +progressbar>=2.3 \ No newline at end of file diff --git a/script.sh b/script.sh old mode 100644 new mode 100755 index 54ab420..2e92428 --- a/script.sh +++ b/script.sh @@ -1,3 +1,4 @@ +#!/usr/bin/env bash echo "Preprocessing" python3 src/preproc.py $1 @@ -5,4 +6,4 @@ echo "Optimizing trajectory" python3 src/stabilize.py $2 echo "Generating output" -python3 src/generate.py $1 $2 \ No newline at end of file +python3 src/generate.py $1 $2 diff --git a/src/preproc.py b/src/preproc.py index 18e4c90..723a289 100644 --- a/src/preproc.py +++ b/src/preproc.py @@ -2,8 +2,11 @@ import numpy as np import pickle import sys +from progressbar import ProgressBar, Percentage, Bar, ETA sift = cv2.xfeatures2d.SIFT_create(200) +pbar = ProgressBar(widgets=[Percentage(), Bar(), ETA()]) + def getAffMat(I1, I2): I1 = cv2.cvtColor(I1, cv2.COLOR_BGR2GRAY) @@ -16,17 +19,17 @@ def getAffMat(I1, I2): # Finding good matches using ratio testing bf = cv2.BFMatcher() matches = bf.knnMatch(desc1, desc2, k=2) - + good = [] - for m,n in matches: + for m, n in matches: if m.distance < 0.7*n.distance: good.append(m) pts_src = [] pts_dst = [] for i in range(len(good)): - pts_src.append([kp1[good[i].queryIdx].pt[0], kp1[good[i].queryIdx].pt[1]]) - pts_dst.append([kp2[good[i].trainIdx].pt[0], kp2[good[i].trainIdx].pt[1]]) + pts_src.append([kp1[good[i].queryIdx].pt[0], kp1[good[i].queryIdx].pt[1]]) + pts_dst.append([kp2[good[i].trainIdx].pt[0], kp2[good[i].trainIdx].pt[1]]) pts_src = np.array(pts_src).astype(np.float32) pts_dst = np.array(pts_dst).astype(np.float32) @@ -34,14 +37,17 @@ def getAffMat(I1, I2): # Computing affine matrix using the best matches return cv2.estimateAffinePartial2D(pts_src, pts_dst)[0] + v = cv2.VideoCapture(sys.argv[1]) n_frames = int(v.get(cv2.CAP_PROP_FRAME_COUNT)) +pbar.maxval = n_frames +pbar.start() # Generating the Xdata and Ydata transforms = [[], [], [], []] count = 0 while v.isOpened(): - + ret, frame = v.read() if ret == True: @@ -54,19 +60,19 @@ def getAffMat(I1, I2): transforms[2].append(np.arctan2(transMat[1][0], transMat[0][0])) transforms[3].append(np.sqrt(transMat[1][0]**2 + transMat[0][0]**2)) except: - transforms[0].append(0) - transforms[1].append(0) - transforms[2].append(0) - transforms[3].append(1) - + transforms[0].append(0) + transforms[1].append(0) + transforms[2].append(0) + transforms[3].append(1) + count += 1 prev = frame - print( str((count/n_frames)*100) + "% completed") + pbar.update(count) else: - break + break v.release() # Storing the data with open('transforms.pkl', 'wb') as f: - pickle.dump(transforms, f) \ No newline at end of file + pickle.dump(transforms, f)