U !?h @s\ddlZddlmZddlZddlZddlmZddlmZ ddl m Z Gddde Z dS)N)List)Image) BaseSessionc@s<eZdZdZeeedddZeddZeddZ d S) U2netSessionzT This class represents a U2net session, which is a subclass of BaseSession. )imgreturnc Os|jd||ddd}|ddddddddf}t|}t|}||||}t|}tj|d ddd }| |j tj j }|gS) aJ Predicts the output masks for the input image using the inner session. Parameters: img (PILImage): The input image. *args: Additional positional arguments. **kwargs: Additional keyword arguments. Returns: List[PILImage]: The list of output masks. N)g ףp= ?gv/?gCl?)gZd;O?gy&1?g?)@r rZuint8L)mode)Z inner_sessionrun normalizenpmaxminZsqueezerZ fromarrayZastyperesizesizeZ ResamplingZLANCZOS) selfrargskwargsZort_outspredmamimaskrF/var/www/html/venv/lib/python3.8/site-packages/rembg/sessions/u2net.pypredicts"     zU2netSession.predictcOsT|j||d}tjd|j||r(dnd||j||ddtj|j|||S)a Downloads the U2net model file from a specific URL and saves it. Parameters: *args: Additional positional arguments. **kwargs: Additional keyword arguments. Returns: str: The path to the downloaded model file. z.onnxzHhttps://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnxNz$md5:60024c5c889badc19c04ad937298a77bT)fnamepathZ progressbar)namepoochretrieveZchecksum_disabledZ u2net_homeosrjoin)clsrrrrrrdownload_models1s    zU2netSession.download_modelscOsdS)z Returns the name of the U2net session. Parameters: *args: Additional positional arguments. **kwargs: Additional keyword arguments. Returns: str: The name of the session. Zu2netr)r%rrrrrr Ls zU2netSession.nameN) __name__ __module__ __qualname____doc__PILImagerr classmethodr&r rrrrr s   r) r#typingrnumpyrr!ZPILrZ PIL.Imager+baserrrrrrs