David310/SAM2_video_mask_generator

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3 4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7from setuptools import find_packages, setup8from torch.utils.cpp_extension import BuildExtension, CUDAExtension9 10# Package metadata11NAME = "SAM 2"12VERSION = "1.0"13DESCRIPTION = "SAM 2: Segment Anything in Images and Videos"14URL = "https://github.com/facebookresearch/segment-anything-2"15AUTHOR = "Meta AI"16AUTHOR_EMAIL = "segment-anything@meta.com"17LICENSE = "Apache 2.0"18 19# Read the contents of README file20with open("README.md", "r") as f:21    LONG_DESCRIPTION = f.read()22 23# Required dependencies24REQUIRED_PACKAGES = [25    "torch>=2.3.1",26    "torchvision>=0.18.1",27    "numpy>=1.24.4",28    "tqdm>=4.66.1",29    "hydra-core>=1.3.2",30    "iopath>=0.1.10",31    "pillow>=9.4.0",32]33 34EXTRA_PACKAGES = {35    "demo": ["matplotlib>=3.9.1", "jupyter>=1.0.0", "opencv-python>=4.7.0"],36    "dev": ["black==24.2.0", "usort==1.0.2", "ufmt==2.0.0b2"],37}38 39 40def get_extensions():41    srcs = ["sam2/csrc/connected_components.cu"]42    compile_args = {43        "cxx": [],44        "nvcc": [45            "-DCUDA_HAS_FP16=1",46            "-D__CUDA_NO_HALF_OPERATORS__",47            "-D__CUDA_NO_HALF_CONVERSIONS__",48            "-D__CUDA_NO_HALF2_OPERATORS__",49        ],50    }51    ext_modules = [CUDAExtension("sam2._C", srcs, extra_compile_args=compile_args)]52    return ext_modules53 54 55# Setup configuration56setup(57    name=NAME,58    version=VERSION,59    description=DESCRIPTION,60    long_description=LONG_DESCRIPTION,61    long_description_content_type="text/markdown",62    url=URL,63    author=AUTHOR,64    author_email=AUTHOR_EMAIL,65    license=LICENSE,66    packages=find_packages(exclude="notebooks"),67    install_requires=REQUIRED_PACKAGES,68    extras_require=EXTRA_PACKAGES,69    python_requires=">=3.10.0",70    ext_modules=get_extensions(),71    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},72)73