Merge pull request #6 from pypose/fix-imu-source

Fix imu source

Author: wang-chen

imu_source/imu_corrector_tutorial.py

@@ -35,6 +35,7 @@
 import matplotlib.pyplot as plt
 from matplotlib.patches import Ellipse
 from matplotlib.collections import PatchCollection
+from torchvision.datasets.utils import download_and_extract_archive
 
 ######################################################################
 # 1. Dataset Defination
@@ -61,8 +62,12 @@
 # 
 
 class KITTI_IMU(Data.Dataset):
-    def __init__(self, root, dataname, drive, duration=10, step_size=1, mode='train'):
+    datalink = 'https://github.com/pypose/pypose/releases/download/v0.2.2/2011_09_26.zip'
+    def __init__(self, root, dataname, drive, duration=10, step_size=1, mode='train',
+                 download=True):
         super().__init__()
+        if download:
+            download_and_extract_archive(self.datalink, root)
         self.duration = duration
         self.data = pykitti.raw(root, dataname, drive)
         self.seq_len = len(self.data.timestamps) - 1

imu_source/imu_dataset_tutorial.py

@@ -10,11 +10,15 @@
 import pypose as pp
 from datetime import datetime
 import torch.utils.data as Data
-
+from torchvision.datasets.utils import download_and_extract_archive
 
 class KITTI_IMU(Data.Dataset):
-    def __init__(self, root, dataname, drive, duration=10, step_size=1, mode='train'):
+    datalink = 'https://github.com/pypose/pypose/releases/download/v0.2.2/2011_09_26.zip'
+    def __init__(self, root, dataname, drive, duration=10, step_size=1, mode='train',
+                 download=True):
         super().__init__()
+        if download:
+            download_and_extract_archive(self.datalink, root)
         self.duration = duration
         self.data = pykitti.raw(root, dataname, drive)
         self.seq_len = len(self.data.timestamps) - 1

imu_source/imu_integrator_tutorial.py

@@ -51,6 +51,7 @@
 import matplotlib.pyplot as plt
 from matplotlib.patches import Ellipse
 from matplotlib.collections import PatchCollection
+from torchvision.datasets.utils import download_and_extract_archive
 
 
 ######################################################################
@@ -78,8 +79,12 @@
 # 
 
 class KITTI_IMU(Data.Dataset):
-    def __init__(self, root, dataname, drive, duration=10, step_size=1, mode='train'):
+    datalink = 'https://github.com/pypose/pypose/releases/download/v0.2.2/2011_09_26.zip'
+    def __init__(self, root, dataname, drive, duration=10, step_size=1, mode='train',
+                 download=True):
         super().__init__()
+        if download:
+            download_and_extract_archive(self.datalink, root)
         self.duration = duration
         self.data = pykitti.raw(root, dataname, drive)
         self.seq_len = len(self.data.timestamps) - 1