[fix] Fix LMDBFeatureReader (#116)

Summary:
Fix LMDBFeatureReader

Test Plan : Test with mmimdb dataset.
Pull Request resolved: https://github.com/fairinternal/pythia-internal/pull/116

Reviewed By: apsdehal

Differential Revision: D21445705

Pulled By: vedanuj

fbshipit-source-id: af96f3cc32f50d3c97ebc2df8ca815f97bea0c0d

Author: vedanuj

mmf/datasets/databases/readers/feature_readers.py

@@ -202,16 +202,14 @@ def _load(self, image_file_path):
             image_id = int(split.split("_")[-1])
             # Try fetching to see if it actually exists otherwise fall back to
             # default
-            _ = self.image_id_indices[image_id]
+            img_id_idx = self.image_id_indices[str(image_id).encode()]
         except (ValueError, KeyError):
             # The image id is complex or involves folder, use it directly
-            image_id = split
-        image_id = str(image_id).encode()
+            image_id = str(split).encode()
+            img_id_idx = self.image_id_indices[image_id]
 
         with self.env.begin(write=False, buffers=True) as txn:
-            image_info = pickle.loads(
-                txn.get(self.image_ids[self.image_id_indices[image_id]])
-            )
+            image_info = pickle.loads(txn.get(self.image_ids[img_id_idx]))
 
         return image_info