diff --git a/datasets.py b/datasets.py
index 0d4e239f97aa509eb4b547ea112d0adc2c1e4104..f04f86d52bf52b12e624d2a95c49294f7a9340e8 100755
--- a/datasets.py
+++ b/datasets.py
@@ -33,8 +33,8 @@ class GenericDataset:
         self.name = name
         self.data_arr = data_arr
         if self.name == 'KITTI':
-            #with open(r"/root/lost/datasets/kitti_labels.pkl", "rb") as input_file:
-            with open(r"/root/lost/Kitti2Coco/train/kitti_labels.pkl", "rb") as input_file:
+            # TODO need to not hardcode
+            with open(r"tools/Kitti2Coco/kitti_train_labels.pkl", "rb") as input_file:
                 self.annots = pickle.load(input_file)
                 print(len(self.data_arr))
 
@@ -61,8 +61,6 @@ class GenericDataset:
                 return [img, self.data_arr[i]]
             if self.annots != None:
                 if self.name == 'KITTI':
-                    print(self.annots[im_name]['annotations'])
-                    print(self.annots[im_name])
                     return [img, self.data_arr[i], self.annots[im_name]['annotations'], img.size, self.annots[im_name]]
                 return [img, self.data_arr[i], self.annots[im_name], img.size]
 
@@ -74,8 +72,14 @@ class GenericDataset:
             return None
         
         if self.name == 'KITTI':
-            return None # TODO need to handle returning annotations
-
+            annots = self.annots[im_name]['annotations']
+            gt_bbxs = []
+            gt_clss = []
+            for gt in annots:
+                gt_bbxs.append(gt['bbox'])
+                gt_clss.append(gt['category_id'])
+            return np.asarray(gt_bbxs), gt_clss
+        
         im = self.annots[im_name]
         # {"labels": ['bbox_x1','bbox_y1','bbox_x2','bbox_y2','class', 'test']}
         gt_bbxs = im[0:4]
diff --git a/main_lost.py b/main_lost.py
index 967395dfc887cf934ef6844bd01bf991e2f9e06c..b7450b268b8867a1cdc56c4212db4319424b0411 100755
--- a/main_lost.py
+++ b/main_lost.py
@@ -128,7 +128,7 @@ if __name__ == "__main__":
     if args.image_path is not None:
         dataset = ImageDataset(args.image_path)
     elif args.dataset == "KITTI":
-        dataset = ImageFolderDataset("$KITTI_ROOT/training/image_2/") # TODO dont hard code
+        dataset = ImageFolderDataset('KITTI', os.environ.get('$KITTI_ROOT','/root/kitti')+'/training/image_2/') # TODO dont hard code
     else:
         dataset = Dataset(args.dataset, args.set, args.no_hard)
 
diff --git a/scripts/run-dataset.sh b/scripts/run-dataset.sh
index 3566356da0f7c5b6cf49f7890128236eb097749e..fa2a485210d9db7118f499feac02793471a40c84 100755
--- a/scripts/run-dataset.sh
+++ b/scripts/run-dataset.sh
@@ -1,5 +1,5 @@
 
-OUTPUT_PATH=/root/kitti/lost_output
+OUTPUT_PATH=/root/lost/outputs/kitti
 
 DINO_ARCH=vit_base
 LOST_FEATURES=k
diff --git a/visualizations.py b/visualizations.py
index a4ec6990279cc8959be1069ae1c64b57de44c243..0ef400940baf1cf1d9b440aac6b814535171e699 100755
--- a/visualizations.py
+++ b/visualizations.py
@@ -43,7 +43,7 @@ def visualize_predictions(image, pred, seed, scales, dims, vis_folder, im_name,
             (int(pred[2]), int(pred[3])),
             (0, 255, 0), 3,
         )
-    print("image.shape:",image.shape, "\npred_box: [x1,y1,x2,y2]", pred)
+    #print("image.shape:",image.shape, "\npred_box: [x1,y1,x2,y2]", pred)
     # Plot the seed
     if plot_seed:
         s_ = np.unravel_index(seed.cpu().numpy(), (w_featmap, h_featmap))