Fix error with exporting quantized saved model when labels cache not …

…present (ultralytics#4174) Co-authored-by: Glenn Jocher <[email protected]>
CV4EcologySchool · Aug 5, 2023 · 730af83 · 730af83
1 parent f0ad1a7
commit 730af83
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/ultralytics/data/annotator.py b/ultralytics/data/annotator.py
@@ -24,13 +24,13 @@ def auto_annotate(data, det_model='yolov8x.pt', sam_model='sam_b.pt', device='',
     det_results = det_model(data, stream=True, device=device)
 
     for result in det_results:
-        boxes = result.boxes.xyxy  # Boxes object for bbox outputs
         class_ids = result.boxes.cls.int().tolist()  # noqa
         if len(class_ids):
+            boxes = result.boxes.xyxy  # Boxes object for bbox outputs
             sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
             segments = sam_results[0].masks.xyn  # noqa
 
-            with open(str(Path(output_dir) / Path(result.path).stem) + '.txt', 'w') as f:
+            with open(f'{str(Path(output_dir) / Path(result.path).stem)}.txt', 'w') as f:
                 for i in range(len(segments)):
                     s = segments[i]
                     if len(s) == 0:

diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py
@@ -605,7 +605,8 @@ def export_saved_model(self, prefix=colorstr('TensorFlow SavedModel:')):
 
                 # Generate calibration data for integer quantization
                 LOGGER.info(f"{prefix} collecting INT8 calibration images from 'data={self.args.data}'")
-                dataset = YOLODataset(check_det_dataset(self.args.data)['val'], imgsz=self.imgsz[0], augment=False)
+                data = check_det_dataset(self.args.data)
+                dataset = YOLODataset(data['val'], data=data, imgsz=self.imgsz[0], augment=False)
                 images = []
                 n_images = 100  # maximum number of images
                 for n, batch in enumerate(dataset):