diff --git a/MTSP-main/ultralytics/models/fastsam/val.py b/MTSP-main/ultralytics/models/fastsam/val.py deleted file mode 100644 index 52e6889..0000000 --- a/MTSP-main/ultralytics/models/fastsam/val.py +++ /dev/null @@ -1,232 +0,0 @@ -# Ultralytics YOLO 🚀, AGPL-3.0 license - -from multiprocessing.pool import ThreadPool -from pathlib import Path - -import numpy as np -import torch -import torch.nn.functional as F - -from ultralytics.models.yolo.detect import DetectionValidator -from ultralytics.utils import LOGGER, NUM_THREADS, ops -from ultralytics.utils.checks import check_requirements -from ultralytics.utils.metrics import SegmentMetrics, box_iou, mask_iou -from ultralytics.utils.plotting import output_to_target, plot_images - -# FastSAMValidator类继承自DetectionValidator,主要用于对FastSAM模型在分割任务中的验证相关操作,例如处理预测结果、计算评估指标、绘制验证图像等 -class FastSAMValidator(DetectionValidator): - - def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None): - """ - Initialize SegmentationValidator and set task to 'segment', metrics to SegmentMetrics. - - 函数功能: - 初始化FastSAMValidator类,调用父类(DetectionValidator)的初始化方法,同时设置任务类型为'segment', - 并初始化用于评估分割任务的指标对象(SegmentMetrics)。 - - 参数说明: - dataloader (可选):数据加载器,用于加载验证数据。 - save_dir (可选):保存验证结果的目录路径。 - pbar (可选):进度条对象,用于显示验证进度(可能在可视化进度相关的功能中使用)。 - args (可选):包含各种配置参数的对象,例如模型相关的参数、验证相关的设置等。 - _callbacks (可选):回调函数相关对象,用于在特定事件发生时执行自定义的操作(比如在验证过程的某些阶段触发额外的处理逻辑)。 - """ - super().__init__(dataloader, save_dir, pbar, args, _callbacks) - self.args.task = 'segment' - self.metrics = SegmentMetrics(save_dir=self.save_dir, on_plot=self.on_plot) - - def preprocess(self, batch): - """ - Preprocesses batch by converting masks to float and sending to device. - - 函数功能: - 对输入的批次数据(batch)进行预处理,先调用父类的预处理方法,然后将批次数据中的掩码(masks)数据转换为浮点数类型,并发送到指定的设备(self.device)上。 - - 参数说明: - batch:包含了图像、标签、掩码等多种数据的批次数据,格式通常是按照数据加载器的定义组织的。 - - 返回值: - 处理后的批次数据,其中掩码数据已转换为浮点数并放置在相应设备上。 - """ - batch = super().preprocess(batch) - batch['masks'] = batch['masks'].to(self.device).float() - return batch - - def init_metrics(self, model): - """ - Initialize metrics and select mask processing function based on save_json flag. - - 函数功能: - 初始化评估指标相关的设置,先调用父类的初始化指标方法,然后根据是否保存JSON格式结果(self.args.save_json)的配置, - 选择合适的掩码处理函数,用于后续对预测掩码的处理操作。 - - 参数说明: - model:正在验证的模型对象,可能在某些与模型相关的指标初始化操作中会用到(虽然此处代码中未体现具体使用情况)。 - """ - super().init_metrics(model) - self.plot_masks = [] - if self.args.save_json: - check_requirements('pycocotools>=2.0.6') - self.process = ops.process_mask_upsample # more accurate - else: - self.process = ops.process_mask # faster - - def get_desc(self): - - return ('%22s' + '%11s' * 10) % ('Class', 'Images', 'Instances', 'Box(P', 'R', 'mAP50', 'mAP50-95)', 'Mask(P', - 'R', 'mAP50', 'mAP50-95)') - - def postprocess(self, preds): - - p = ops.non_max_suppression(preds[0], - self.args.conf, - self.args.iou, - labels=self.lb, - multi_label=True, - agnostic=self.args.single_cls, - max_det=self.args.max_det, - nc=self.nc) - proto = preds[1][-1] if len(preds[1]) == 3 else preds[1] # second output is len 3 if pt, but only 1 if exported - return p, proto - - def update_metrics(self, preds, batch): - """Metrics.""" - for si, (pred, proto) in enumerate(zip(preds[0], preds[1])): - idx = batch['batch_idx'] == si - cls = batch['cls'][idx] - bbox = batch['bboxes'][idx] - nl, npr = cls.shape[0], pred.shape[0] # number of labels, predictions - shape = batch['ori_shape'][si] - correct_masks = torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device) # init - correct_bboxes = torch.zeros(npr, self.niou, dtype=torch.bool, device=self.device) # init - self.seen += 1 - - if npr == 0: - if nl: - self.stats.append((correct_bboxes, correct_masks, *torch.zeros( - (2, 0), device=self.device), cls.squeeze(-1))) - if self.args.plots: - self.confusion_matrix.process_batch(detections=None, labels=cls.squeeze(-1)) - continue - - # Masks - midx = [si] if self.args.overlap_mask else idx - gt_masks = batch['masks'][midx] - pred_masks = self.process(proto, pred[:, 6:], pred[:, :4], shape=batch['img'][si].shape[1:]) - - # Predictions - if self.args.single_cls: - pred[:, 5] = 0 - predn = pred.clone() - ops.scale_boxes(batch['img'][si].shape[1:], predn[:, :4], shape, - ratio_pad=batch['ratio_pad'][si]) # native-space pred - - # Evaluate - if nl: - height, width = batch['img'].shape[2:] - tbox = ops.xywh2xyxy(bbox) * torch.tensor( - (width, height, width, height), device=self.device) # target boxes - ops.scale_boxes(batch['img'][si].shape[1:], tbox, shape, - ratio_pad=batch['ratio_pad'][si]) # native-space labels - labelsn = torch.cat((cls, tbox), 1) # native-space labels - correct_bboxes = self._process_batch(predn, labelsn) - # TODO: maybe remove these `self.` arguments as they already are member variable - correct_masks = self._process_batch(predn, - labelsn, - pred_masks, - gt_masks, - overlap=self.args.overlap_mask, - masks=True) - if self.args.plots: - self.confusion_matrix.process_batch(predn, labelsn) - - # Append correct_masks, correct_boxes, pconf, pcls, tcls - self.stats.append((correct_bboxes, correct_masks, pred[:, 4], pred[:, 5], cls.squeeze(-1))) - - pred_masks = torch.as_tensor(pred_masks, dtype=torch.uint8) - if self.args.plots and self.batch_i < 3: - self.plot_masks.append(pred_masks[:15].cpu()) # filter top 15 to plot - - # Save - if self.args.save_json: - pred_masks = ops.scale_image(pred_masks.permute(1, 2, 0).contiguous().cpu().numpy(), - shape, - ratio_pad=batch['ratio_pad'][si]) - self.pred_to_json(predn, batch['im_file'][si], pred_masks) - # if self.args.save_txt: - # save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt') - - def finalize_metrics(self, *args, **kwargs): - - self.metrics.speed = self.speed - self.metrics.confusion_matrix = self.confusion_matrix - - def _process_batch(self, detections, labels, pred_masks=None, gt_masks=None, overlap=False, masks=False): - - if masks: - if overlap: - nl = len(labels) - index = torch.arange(nl, device=gt_masks.device).view(nl, 1, 1) + 1 - gt_masks = gt_masks.repeat(nl, 1, 1) # shape(1,640,640) -> (n,640,640) - gt_masks = torch.where(gt_masks == index, 1.0, 0.0) - if gt_masks.shape[1:]!= pred_masks.shape[1:]: - gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode='bilinear', align_corners=False)[0] - gt_masks = gt_masks.gt_(0.5) - iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1)) - else: # boxes - iou = box_iou(labels[:, 1:], detections[:, :4]) - - correct = np.zeros((detections.shape[0], self.iouv.shape[0])).astype(bool) - correct_class = labels[:, 0:1] == detections[:, 5] - for i in range(len(self.iouv)): - x = torch.where((iou >= self.iouv[i]) & correct_class) # IoU > threshold and classes match - if x[0].shape[0]: - matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), - 1).cpu().numpy() # [label, detect, iou] - if x[0].shape[0] > 1: - matches = matches[matches[:, 2].argsort()[::-1]] - matches = matches[np.unique(matches[:, 1], return_index=True)[1]] - # matches = matches[matches[:, 2].argsort()[::-1]] - matches = matches[np.unique(matches[:, 0], return_index=True)[1]] - correct[matches[:, 1].astype(int), i] = True - return torch.tensor(correct, dtype=torch.bool, device=detections.device) - - def plot_val_samples(self, batch, ni): - - plot_images(batch['img'], - batch['batch_idx'], - batch['cls'].squeeze(-1), - batch['bboxes'], - batch['masks'], - paths=batch['im_file'], - fname=self.save_dir / f'val_batch{ni}_labels.jpg', - names=self.names, - on_plot=self.on_plot) - - def plot_predictions(self, batch, preds, ni): - - plot_images( - batch['img'], - *output_to_target(preds[0], max_det=15), # not set to self.args.max_det due to slow plotting speed - torch.cat(self.plot_masks, dim=0) if len(self.plot_masks) else self.plot_masks, - paths=batch['im_file'], - fname=self.save_dir / f'val_batch{ni}_pred.jpg', - names=self.names, - on_plot=self.on_plot) # pred - self.plot_masks.clear() - - def pred_to_json(self, predn, filename, pred_masks): - - from pycocotools.mask import encode # noqa - - def single_encode(x): - - rle = encode(np.asarray(x[:, :, None], order='F', dtype='uint8'))[0] - rle['counts'] = rle['counts'].decode('utf-8') - return rle - - stem = Path(filename).stem - image_id = int(stem) if stem.isnumeric() else stem - box = ops.xyxy2xywh(predn[:, :4]) # xywh - box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner - pred_masks = np.transpose(pred_masks, (2, \ No newline at end of file diff --git a/MTSP-main/ultralytics/models/nas/model.py b/MTSP-main/ultralytics/models/nas/model.py index 1f7cd35..3eb3820 100644 --- a/MTSP-main/ultralytics/models/nas/model.py +++ b/MTSP-main/ultralytics/models/nas/model.py @@ -1,13 +1,3 @@ -# Ultralytics YOLO 🚀, AGPL-3.0 license -""" -YOLO-NAS model interface. - -Usage - Predict: - from ultralytics import NAS - - model = NAS('yolo_nas_s') - results = model.predict('ultralytics/assets/bus.jpg') -""" from pathlib import Path @@ -28,30 +18,23 @@ class NAS(Model): @smart_inference_mode() def _load(self, weights: str, task: str): - # Load or create new NAS model + import super_gradients suffix = Path(weights).suffix if suffix == '.pt': self.model = torch.load(weights) elif suffix == '': self.model = super_gradients.training.models.get(weights, pretrained_weights='coco') - # Standardize model + self.model.fuse = lambda verbose=True: self.model self.model.stride = torch.tensor([32]) self.model.names = dict(enumerate(self.model._class_names)) self.model.is_fused = lambda: False # for info() - self.model.yaml = {} # for info() - self.model.pt_path = weights # for export() - self.model.task = 'detect' # for export() + self.model.yaml = {} + self.model.task = 'detect' def info(self, detailed=False, verbose=True): - """ - Logs model info. - Args: - detailed (bool): Show detailed information about model. - verbose (bool): Controls verbosity. - """ return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640) @property