xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Return processed outputs from predictor (#161)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
This commit is contained in:
Ayush Chaurasia
2023-01-10 00:10:44 +05:30
committed by GitHub
parent cb4801888e
commit 6e5638c128
7 changed files with 23 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
ultralytics/yolo/v8/segment/predict.py
View File

@ -58,7 +58,7 @@ class SegmentationPredictor(DetectionPredictor):
mask = masks[idx]
if self.args.save_txt:
segments = [
ops.scale_segments(im0.shape if self.arg.retina_masks else im.shape[2:], x, im0.shape, normalize=True)
ops.scale_segments(im0.shape if self.args.retina_masks else im.shape[2:], x, im0.shape, normalize=True)
for x in reversed(ops.masks2segments(mask))]
# Print results
@ -73,6 +73,9 @@ class SegmentationPredictor(DetectionPredictor):
im_gpu=torch.as_tensor(im0, dtype=torch.float16).to(self.device).permute(2, 0, 1).flip(0).contiguous() /
255 if self.args.retina_masks else im[idx])
det = reversed(det[:, :6])
self.all_outputs.append([det, mask])
# Write results
for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
if self.args.save_txt: # Write to file
@ -96,7 +99,7 @@ class SegmentationPredictor(DetectionPredictor):
@hydra.main(version_base=None, config_path=str(DEFAULT_CONFIG.parent), config_name=DEFAULT_CONFIG.name)
def predict(cfg):
cfg.model = cfg.model or "n.pt"
cfg.model = cfg.model or "yolov8n-seg.pt"
cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2) # check image size
predictor = SegmentationPredictor(cfg)
predictor()