https://github.com/ria-com/nomeroff-net/issues/280

ria-com · Dec 16, 2024 · fb1cd69 · fb1cd69
1 parent 10af341
commit fb1cd69
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/nomeroff_net/tools/datasets_tools.py b/nomeroff_net/tools/datasets_tools.py
@@ -152,9 +152,13 @@ def auto_number_grab(root_dir, res_dir, replace_template=None, csv_dataset_path=
     number_plate_detection_and_reading = pipeline("number_plate_detection_and_reading", image_loader=image_loader,
                                                   **kwargs)
 
-    photos = pd.DataFrame(columns=['photoId'])
+    photos = pd.DataFrame(columns=['photoId', 'npText'])
     if csv_dataset_path is not None:
-        photos = pd.read_csv(csv_dataset_path)
+        photos = pd.read_csv(
+            csv_dataset_path,
+            usecols=['photoId', 'npText'],
+            dtype={'photoId': int, 'npText': str}
+        )
         photos = photos.set_index(['photoId'])
 
     if os.path.exists(res_dir):