ImageDatasetAnalyzer/setup.py at main · joortif/ImageDatasetAnalyzer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
from setuptools import setup, find_packages
from codecs import open

# For installing PyTorch and Torchvision in Windows
import sys
import subprocess
import pkg_resources

with open('README.md', 'r', encoding='utf-8') as f:
    long_description = f.read()

def parse_requirements(filename):
    """ load requirements from a pip requirements file """
    lineiter = (line.strip() for line in open(filename))
    return [line for line in lineiter if line and not line.startswith("#")]

install_reqs = parse_requirements("requirements.txt")

def remove_requirements(requirements, remove_elem):
    new_requirements = []
    for requirement in requirements:
        if remove_elem not in requirement:
            new_requirements.append(requirement)
    return new_requirements

def install_pytorch_for_gpu():
    """Try to install PyTorch and Torchvision with GPU support (CUDA)"""
    print('Checking for GPU support...')
    cuda_version = "cu118"  # Update this if you want to use a different CUDA version
    torch_version = f"torch==2.1.0+{cuda_version}"
    torchvision_version = f"torchvision==0.16.0+{cuda_version}"

    try:
        subprocess.check_call([
            'pip', 'install', torch_version, torchvision_version,
            '-f', 'https://download.pytorch.org/whl/torch_stable.html'
        ])
    except subprocess.CalledProcessError as e:
        print(f"Error installing PyTorch GPU version: {e}")
        print("Trying to install the CPU version instead...")
        install_pytorch_for_cpu()

def install_pytorch_for_cpu():
    """Fallback to install PyTorch and Torchvision with CPU support"""
    torch_version = "torch==2.1.0"
    torchvision_version = "torchvision==0.16.0"

    try:
        subprocess.check_call([
            'pip', 'install', torch_version, torchvision_version,
            '-f', 'https://download.pytorch.org/whl/torch_stable.html'
        ])
        print('Successfully installed PyTorch and Torchvision for CPU.')
    except subprocess.CalledProcessError as e:
        print(f"Error installing PyTorch CPU version: {e}")
        print("Please install PyTorch and Torchvision manually following the instructions at: https://pytorch.org/get-started/")

# Remove PyTorch and Torchvision from install requirements to handle manually
install_reqs = remove_requirements(install_reqs, 'torch')
install_reqs = remove_requirements(install_reqs, 'torchvision')

def check_and_install_pytorch():
    """Check if PyTorch and Torchvision are installed and compatible"""
    try:
        # Check if torch and torchvision are installed and compatible with the versions in requirements.txt
        pkg_resources.require("torch>=2.1.0")
        pkg_resources.require("torchvision>=0.16.0")
        print("Required versions of PyTorch and Torchvision are already installed.")
    except (pkg_resources.DistributionNotFound, pkg_resources.VersionConflict):
        print("PyTorch or Torchvision not found or incompatible, installing the correct versions...")
        install_pytorch_for_gpu()

print(f"Platform: {sys.platform}")
if sys.platform in ['win32', 'cygwin', 'windows']:
    # Check and install the correct versions of PyTorch and Torchvision if necessary
    check_and_install_pytorch()

setup(
    name = 'imagedatasetanalyzer',

    version = '0.4.3',

    author = 'Joaquin Ortiz de Murua Ferrero',
    author_email = 'jortizdemuruaferrero@gmail.com',
    maintainer= 'Joaquin Ortiz de Murua Ferrero',
    maintainer_email= 'jortizdemuruaferrero@gmail.com',

    url='https://github.com/joortif/ImageDatasetAnalyzer',

    description = 'Image dataset analyzer using image embedding models and clustering methods.',

    long_description_content_type = 'text/markdown',
    long_description = long_description,

    license = 'MIT license',

    packages = find_packages(exclude=["test"]),
    install_requires = install_reqs,
    include_package_data=True,

    classifiers=[

        'Development Status :: 4 - Beta',

        'Programming Language :: Python :: 3.10',

        'Intended Audience :: Science/Research',
        'Intended Audience :: Developers',

        # Topics
        'Topic :: Scientific/Engineering :: Image Recognition',
        'Topic :: Scientific/Engineering :: Artificial Intelligence',

        "Operating System :: OS Independent",
    ],

    keywords='instance semantic segmentation pytorch tensorflow huggingface opencv embedding image analysis machine learning deep learning active learning computer vision'
)