
    j                    V    d dl mZ d dlmZ d dlmZ d dlZd dlmZ 	 	 	 	 	 	 	 dddZ	dS )    )annotations)Path)AnyN)LOGGERF modeltorch.nn.Moduleim<torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...]
output_dirPath | str | Nonedynamicboolhalfint8calibration_dataset
Any | Noneignored_scopedict | Noneprefixstrreturnr   c	                   ddl }	t          j        d| d|	j         d           t	          |t
          t          f          rd |D             n|j        }
|	                    | |rdn|
|          }|r'ddl	}|
                    |||j        j        |          }|Ct          |          }|                    d	d	
           |dz  }|	                    |||           |S )aQ  Export a PyTorch model to OpenVINO format with optional INT8 quantization.

    Args:
        model (torch.nn.Module): The model to export (may be NMS-wrapped).
        im (torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...]): Example input tensor(s) for tracing.
        output_dir (Path | str | None): Directory to save the exported OpenVINO model.
        dynamic (bool): Whether to use dynamic input shapes.
        half (bool): Whether to compress to FP16.
        int8 (bool): Whether to apply INT8 quantization.
        calibration_dataset (nncf.Dataset | None): Dataset for INT8 calibration (required when ``int8=True``).
        ignored_scope (dict | None): Kwargs passed to ``nncf.IgnoredScope`` for head patterns.
        prefix (str): Prefix for log messages.

    Returns:
        (ov.Model): The converted OpenVINO model.
    r   N
z starting export with openvino z...c                    g | ]	}|j         
S  )shape).0is     d/var/www/html/Carbon-Document/venv/lib/python3.11/site-packages/ultralytics/utils/export/openvino.py
<listcomp>z"torch2openvino.<locals>.<listcomp>,   s    '''q17'''    )inputexample_input)r   r   presetr   T)parentsexist_okz	model.xml)compress_to_fp16)openvinor   info__version__
isinstancelisttupler   convert_modelnncfquantizeQuantizationPresetMIXEDr   mkdir
save_model)r   r
   r   r   r   r   r   r   r   ovinput_shapeov_modelr0   output_files                 r    torch2openvinor:      s   6 
KOVOOBNOOOPPP+5b4-+H+HV''B''''bhKW-MTT+]_``H 
== 3*0'	 ! 
 
 *%%
555 ;.
hdCCCOr"   )NFFFNNr   )r   r	   r
   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   )

__future__r   pathlibr   typingr   torchultralytics.utilsr   r:   r   r"   r    <module>r@      s    # " " " " "              $ $ $ $ $ $ %)&*!%0 0 0 0 0 0 0r"   