o V hT@sdZddlZddlZddlZddlmZddlZddlZddl Z ddl m Z m Z ddl mZddlmZmZddlmZddlmZmZmZmZmZmZmZdd lmZmZdd lm Z dd l!m"Z"m#Z#d Z$Gd ddZ%dS)ar Run prediction on images, videos, directories, globs, YouTube, webcam, streams, etc. Usage - sources: $ yolo mode=predict model=yolo11n.pt source=0 # webcam img.jpg # image vid.mp4 # video screen # screenshot path/ # directory list.txt # list of images list.streams # list of streams 'path/*.jpg' # glob 'https://youtu.be/LNwODJXcvt4' # YouTube 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP, TCP stream Usage - formats: $ yolo mode=predict model=yolo11n.pt # PyTorch yolo11n.torchscript # TorchScript yolo11n.onnx # ONNX Runtime or OpenCV DNN with dnn=True yolo11n_openvino_model # OpenVINO yolo11n.engine # TensorRT yolo11n.mlpackage # CoreML (macOS-only) yolo11n_saved_model # TensorFlow SavedModel yolo11n.pb # TensorFlow GraphDef yolo11n.tflite # TensorFlow Lite yolo11n_edgetpu.tflite # TensorFlow Edge TPU yolo11n_paddle_model # PaddlePaddle yolo11n.mnn # MNN yolo11n_ncnn_model # NCNN yolo11n_imx_model # Sony IMX yolo11n_rknn_model # Rockchip RKNN N)Path)get_cfg get_save_dir)load_inference_source) LetterBoxclassify_transforms) AutoBackend) DEFAULT_CFGLOGGERMACOSWINDOWS callbackscolorstrops) check_imgsz check_imshow)increment_path) select_devicesmart_inference_modea inference results will accumulate in RAM unless `stream=True` is passed, causing potential out-of-memory errors for large sources or long-running streams and videos. See https://docs.ultralytics.com/modes/predict/ for help. Example: results = model(source=..., stream=True) # generator of Results objects for r in results: boxes = r.boxes # Boxes object for bbox outputs masks = r.masks # Masks object for segment masks outputs probs = r.probs # Class probabilities for classification outputs c@seZdZdZeddfddZddZddZd d Zd d Z d&ddZ d'ddZ ddZ e d'ddZd(ddZddZd)ddZd*dd Zd!efd"d#Zd!efd$d%ZdS)+ BasePredictora A base class for creating predictors. This class provides the foundation for prediction functionality, handling model setup, inference, and result processing across various input sources. Attributes: args (SimpleNamespace): Configuration for the predictor. save_dir (Path): Directory to save results. done_warmup (bool): Whether the predictor has finished setup. model (torch.nn.Module): Model used for prediction. data (dict): Data configuration. device (torch.device): Device used for prediction. dataset (Dataset): Dataset used for prediction. vid_writer (dict): Dictionary of {save_path: video_writer} for saving video output. plotted_img (numpy.ndarray): Last plotted image. source_type (SimpleNamespace): Type of input source. seen (int): Number of images processed. windows (list): List of window names for visualization. batch (tuple): Current batch data. results (list): Current batch results. transforms (callable): Image transforms for classification. callbacks (dict): Callback functions for different events. txt_path (Path): Path to save text results. _lock (threading.Lock): Lock for thread-safe inference. Methods: preprocess: Prepare input image before inference. inference: Run inference on a given image. postprocess: Process raw predictions into structured results. predict_cli: Run prediction for command line interface. setup_source: Set up input source and inference mode. stream_inference: Stream inference on input source. setup_model: Initialize and configure the model. write_results: Write inference results to files. save_predicted_images: Save prediction visualizations. show: Display results in a window. run_callbacks: Execute registered callbacks for an event. add_callback: Register a new callback function. NcCst|||_t|j|_|jjdurd|j_d|_|jjr$tdd|j_d|_|jj |_ d|_ d|_ d|_ i|_ d|_d|_d|_g|_d|_d|_d|_|pRt|_d|_t|_t|dS)a$ Initialize the BasePredictor class. Args: cfg (str | dict): Path to a configuration file or a configuration dictionary. overrides (dict | None): Configuration overrides. _callbacks (dict | None): Dictionary of callback functions. Ng?FT)warnr)rargsrsave_dirconf done_warmupshowrmodeldataimgszdevicedataset vid_writer plotted_img source_typeseenwindowsbatchresults transformsr get_default_callbackstxt_path threadingLock_lockadd_integration_callbacks)selfcfg overrides _callbacksr3P/var/www/vscode/kcb/lib/python3.10/site-packages/ultralytics/engine/predictor.py__init__ls0    zBasePredictor.__init__cCst|tj }|r0t||}|jddkr!|ddddf}|d}t|}t |}| |j }|j j r>|n|}|rH|d}|S)z Prepares input image before inference. Args: im (torch.Tensor | List(np.ndarray)): Images of shape (N, 3, h, w) for tensor, [(h, w, 3) x N] for list. .N)rr7) isinstancetorchTensornpstack pre_transformshape transposeascontiguousarray from_numpytorrfp16halffloat)r/im not_tensorr3r3r4 preprocesss    zBasePredictor.preprocesscOs^|jjr|jjst|jt|jddjddnd}|j |g|R|jj ||jj d|S)zGRun inference on a given image using the specified model and arguments.rT)mkdirF)augment visualizeembed) rrNr#tensorrrrr&stemrrMrO)r/rIrkwargsrNr3r3r4 inferences$(zBasePredictor.inferencecsdtdd|Ddk}t|j|o#|jjo#|jjp#t|jddo#|jj |jj dfdd|DS) a Pre-transform input image before inference. Args: im (List[np.ndarray]): Images of shape (N, 3, h, w) for tensor, [(h, w, 3) x N] for list. Returns: (List[np.ndarray]): A list of transformed images. cSsh|]}|jqSr3)rA.0xr3r3r4 sz.BasePredictor.pre_transform..r8dynamicF)autostridecsg|]}|dqS))imager3rT letterboxr3r4 sz/BasePredictor.pre_transform..) lenrrrrectrptgetattrimxrZ)r/rI same_shapesr3r\r4r@s zBasePredictor.pre_transformcCs|S)z6Post-process predictions for an image and return them.r3)r/predsimg orig_imgsr3r3r4 postprocessszBasePredictor.postprocessFcOsB||_|r|j||g|Ri|St|j||g|Ri|S)a Perform inference on an image or stream. Args: source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor | None): Source for inference. model (str | Path | torch.nn.Module | None): Model for inference. stream (bool): Whether to stream the inference results. If True, returns a generator. *args (Any): Additional arguments for the inference method. **kwargs (Any): Additional keyword arguments for the inference method. Returns: (List[ultralytics.engine.results.Results] | generator): Results objects or generator of Results objects. )streamstream_inferencelist)r/sourcerrirrRr3r3r4__call__szBasePredictor.__call__cCs|||}|D]}qdS)a> Method used for Command Line Interface (CLI) prediction. This function is designed to run predictions using the CLI. It sets up the source and model, then processes the inputs in a streaming manner. This method ensures that no outputs accumulate in memory by consuming the generator without storing results. Args: source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor | None): Source for inference. model (str | Path | torch.nn.Module | None): Model for inference. Note: Do not modify this function or remove the generator. The generator ensures that no outputs are accumulated in memory, which is critical for preventing memory issues during long-running predictions. N)rj)r/rlrgen_r3r3r4 predict_clis zBasePredictor.predict_clicCst|jj|jjdd|_|jjdkrt|jjdt|jdnd|_t ||jj |jj |jj d|_ |j j|_t|dd sZ|jjsU|jjsUt|j d ksUtt|j d d grZtti|_dS) z Set up source and inference mode. Args: source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor): Source for inference. r9)rZmin_dimclassifyr(rN)rlr& vid_stridebufferriTi video_flagF)rrrrrZtaskrbrr(rr&rs stream_bufferr r#ri screenshotr_anyr warningSTREAM_WARNINGr!)r/rlr3r3r4 setup_sources4       zBasePredictor.setup_sourcec /s,jjr tdjs|ji|dur|njjjj s,jj r=jj r5j dnj j dddj s^jjjjsKjjrMdnjjjjgjRdd_ dgd___tjjd tjjd tjjd f}d jD]_d j\}}}|d |} Wdn1swY|d.j| g|Ri|} jjrt| t j!r| gn| EdH WdqWdn1swY|d "| | |_#Wdn1swYd t$|} t%| D]O} jd7_|dj&d| |dj&d| |d j&d| dj#| _'jjsKjj sKjj sKjj(r]|| )| t*|| | |7<qjjrltd+|dj#EdHqWdn 1swYj,-D]} t| t.j/r| 0qjjrψjrt1fdd|D}tdt2jjjt3jddg| j4d dR|jj sވjj sވjj5rt$t6j 7d}jj rd|dd|dkdj dnd}tdt8dj |ddS)a7 Stream real-time inference on camera feed and save results to file. Args: source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor | None): Source for inference. model (str | Path | torch.nn.Module | None): Model for inference. *args (Any): Additional arguments for the inference method. **kwargs (Any): Additional keyword arguments for the inference method. Yields: (ultralytics.engine.results.Results): Results objects. NlabelsTparentsexist_okr8)rr)ron_predict_starton_predict_batch_startr9on_predict_postprocess_end@@)rKrSrh on_predict_batch_endc3s |] }|jjdVqdS)rN)tr$rTr/r3r4 msz1BasePredictor.stream_inference..zRSpeed: %.1fms preprocess, %.1fms inference, %.1fms postprocess per image at shape chr7z labels/*.txtz labelsz saved to zResults saved to boldon_predict_end)9rverboser infor setup_modelr-r|rlsavesave_txtrrLrwarmupratritonr bsrrr$r%r&rProfiler run_callbacksrKrSrOr;r<r=rhr'r_rangedtspeedr write_resultsrjoinr!valuescv2 VideoWriterreleasetupleminrbrA save_croprkglobr)r/rlrrrR profilerspathsim0srrIrenivrnlr3rr4rjs   "*             ($  >,2zBasePredictor.stream_inferenceTc Csbt|p|jjt|jj|d|jj|jj|jj|jjd|d|_|jj|_|jj |j_|j dS)z Initialize YOLO model with given parameters and set it to evaluation mode. Args: model (str | Path | torch.nn.Module | None): Model to load or use. verbose (bool): Whether to print verbose output. )rT)weightsrdnnrrFr&fuserN) rrrrrrrrGr&rFeval)r/rrr3r3r4rxs  zBasePredictor.setup_modelc Csd}t|jdkr |d}|jjs|jjs|jjr%||d7}|jj}nt d||}|r5t |dnd}|j d|j |jj dkrEdnd ||_|d j|jd d7}|j|}|j |_ |||jd d d7}|jjs{|jjr|j|jj|jj|jj|jj|jjrdn||d|_|jjr|j|jd|jjd|jj r|j |j d|jj d|jjr|t!||jjr|"t!|j |j#||S)aq Write inference results to a file or directory. Args: i (int): Index of the current image in the batch. p (Path): Path to the current image. im (torch.Tensor): Preprocessed image tensor. s (List[str]): List of result strings. Returns: (str): String with result information. r}r7Nz: z frame (\d+)/r8r~r[roz {:g}x{:g} r9rSz.1fms) line_widthboxesrr~im_gpuz.txt) save_confcrops)r file_name)$r_rAr#rifrom_imgrPr countresearchintrrQmoder*formatr'__str__rrrrrplotr show_boxes show_conf show_labels retina_masksr"rrrstrsave_predicted_imagesname) r/rprIrstringframematchresultr3r3r4rs<  ,   zBasePredictor.write_resultsr}rcCs |j}|jjdvrw|jjdkr|jjnd}|dddd}||jvr\|jjr2t|j ddd t r6d nt r:d nd \}}t j tt||t j|||jd|jdfd |j|<|j|||jjrut ||d|dSdSt tt|d|dS)z Save video predictions as mp4 or images as jpg at specified path. Args: save_path (str): Path to save the results. frame (int): Frame number for video mode. >videorir.r8rz_frames/Tr)z.mp4avc1).aviWMV2)rMJPG)filenamefourccfps frameSizez.jpgN)r"r rrsplitr!r save_framesrrLr r rrr with_suffixVideoWriter_fourccrAwriteimwrite)r/ save_pathrrIr frames_pathsuffixrr3r3r4rs&   z#BasePredictor.save_predicted_imagescCs|j}tdkr+||jvr+|j|t|tjtjBt ||j d|j dt ||t |j jdkr>ddSddS)zDisplay an image in a window.Linuxr8rr[i,N)r"platformsystemr%appendr namedWindow WINDOW_NORMALWINDOW_KEEPRATIO resizeWindowrAimshowwaitKeyr r)r/rrIr3r3r4rs  $zBasePredictor.showeventcCs |j|gD]}||qdS)z2Run all registered callbacks for a specific event.N)r get)r/rcallbackr3r3r4rs zBasePredictor.run_callbackscCs|j||dS)z-Add a callback function for a specific event.N)r r)r/rfuncr3r3r4 add_callbackszBasePredictor.add_callback)NNF)NN)T)r}r)r})__name__ __module__ __qualname____doc__r r5rKrSr@rhrmrpr|rrjrrrrrrrr3r3r3r4rBs$)$   "  d 3 " r)&rrrr+pathlibrrnumpyr>r<ultralytics.cfgrrultralytics.datarultralytics.data.augmentrrultralytics.nn.autobackendrultralytics.utilsr r r r r rrultralytics.utils.checksrrultralytics.utils.filesrultralytics.utils.torch_utilsrrr{rr3r3r3r4s$!   $