diff --git a/src/features/preprocess_ki.py b/src/features/preprocess_ki.py
index 72b9765..3016562 100644
--- a/src/features/preprocess_ki.py
+++ b/src/features/preprocess_ki.py
@@ -66,11 +66,11 @@ class PreprocessKitti:
         self.set_val = set_gt.intersection(set_val)
         assert self.set_train and self.set_val, "No validation or training annotations"
 
-        self.dic_jo = {'train': dict(X=[], Y=[], names=[], kps=[],
+        self.dic_jo = {'train': dict(X=[], Y=[], names=[], kps=[], K=[],
                                      clst=defaultdict(lambda: defaultdict(list))),
-                       'val': dict(X=[], Y=[], names=[], kps=[],
+                       'val': dict(X=[], Y=[], names=[], kps=[], K=[],
                                    clst=defaultdict(lambda: defaultdict(list))),
-                       'test': dict(X=[], Y=[], names=[], kps=[],
+                       'test': dict(X=[], Y=[], names=[], kps=[], K=[],
                                     clst=defaultdict(lambda: defaultdict(list)))}
 
         self.dic_names = defaultdict(lambda: defaultdict(list))
diff --git a/src/predict/predict_2d_3d.py b/src/predict/predict_2d_3d.py
index bf0d637..d02e012 100644
--- a/src/predict/predict_2d_3d.py
+++ b/src/predict/predict_2d_3d.py
@@ -1,6 +1,4 @@
-"""Predict poses for given images."""
 
-import argparse
 import glob
 import json
 import os