From 730af8390b021b25a829ce2c575f728e8b025a6e Mon Sep 17 00:00:00 2001 From: Colin Wong Date: Sat, 5 Aug 2023 08:35:14 -0500 Subject: [PATCH] Fix error with exporting quantized saved model when labels cache not present (#4174) Co-authored-by: Glenn Jocher --- ultralytics/data/annotator.py | 4 ++-- ultralytics/engine/exporter.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/ultralytics/data/annotator.py b/ultralytics/data/annotator.py index f69f325b..f8d59019 100644 --- 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 index 05789b57..e312f76a 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -605,7 +605,8 @@ class Exporter: # 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):