435 lines
16 KiB
Python
435 lines
16 KiB
Python
"""
|
|
Fix OpenCLIP GPU Dependencies and Complete Isolation from PaddleOCR
|
|
Create separate virtual environment with GPU-enabled PyTorch for OpenCLIP
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import tempfile
|
|
import platform
|
|
|
|
def check_system_gpu():
|
|
"""Check system GPU capabilities"""
|
|
print("🔍 CHECKING SYSTEM GPU CAPABILITIES")
|
|
print("=" * 40)
|
|
|
|
try:
|
|
# Check NVIDIA GPU
|
|
result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
print("✅ NVIDIA GPU detected")
|
|
# Parse driver version
|
|
for line in result.stdout.split('\n'):
|
|
if 'Driver Version' in line:
|
|
print(f" Driver: {line.strip()}")
|
|
if 'CUDA Version' in line:
|
|
print(f" CUDA: {line.strip()}")
|
|
else:
|
|
print("❌ No NVIDIA GPU detected or nvidia-smi not available")
|
|
return False
|
|
|
|
# Check PyTorch CUDA in main environment
|
|
try:
|
|
import torch
|
|
print(f"✅ Main PyTorch CUDA: {torch.cuda.is_available()}")
|
|
if torch.cuda.is_available():
|
|
print(f" CUDA Version: {torch.version.cuda}")
|
|
print(f" GPU Count: {torch.cuda.device_count()}")
|
|
print(f" Current GPU: {torch.cuda.get_device_name(0)}")
|
|
except Exception as e:
|
|
print(f"⚠️ Main PyTorch check failed: {e}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ GPU check failed: {e}")
|
|
return False
|
|
|
|
def create_openclip_gpu_environment():
|
|
"""Create isolated OpenCLIP environment with GPU support"""
|
|
print("\n🛠️ CREATING OPENCLIP GPU ENVIRONMENT")
|
|
print("=" * 40)
|
|
|
|
venv_path = "openclip_gpu_env"
|
|
|
|
# Remove existing environment if it exists
|
|
if os.path.exists(venv_path):
|
|
print(f"🗑️ Removing existing environment: {venv_path}")
|
|
import shutil
|
|
shutil.rmtree(venv_path)
|
|
|
|
try:
|
|
# Create virtual environment
|
|
print(f"📦 Creating virtual environment: {venv_path}")
|
|
result = subprocess.run([
|
|
sys.executable, '-m', 'venv', venv_path,
|
|
'--system-site-packages' # Use system CUDA libraries
|
|
], capture_output=True, text=True)
|
|
|
|
if result.returncode != 0:
|
|
print(f"❌ Virtual environment creation failed: {result.stderr}")
|
|
return False
|
|
|
|
print("✅ Virtual environment created successfully")
|
|
|
|
# Get virtual environment Python path
|
|
if platform.system() == "Windows":
|
|
venv_python = os.path.join(venv_path, "Scripts", "python.exe")
|
|
else:
|
|
venv_python = os.path.join(venv_path, "bin", "python")
|
|
|
|
# Install GPU-enabled PyTorch
|
|
print("📥 Installing GPU-enabled PyTorch...")
|
|
|
|
# Use PyTorch with CUDA 11.8 (compatible with most systems)
|
|
torch_install_cmd = [
|
|
venv_python, '-m', 'pip', 'install',
|
|
'torch==2.0.1+cu118', 'torchvision==0.15.2+cu118',
|
|
'--index-url', 'https://download.pytorch.org/whl/cu118'
|
|
]
|
|
|
|
result = subprocess.run(torch_install_cmd, capture_output=True, text=True, timeout=300)
|
|
if result.returncode != 0:
|
|
print(f"❌ PyTorch installation failed: {result.stderr}")
|
|
return False
|
|
|
|
print("✅ PyTorch with CUDA 11.8 installed")
|
|
|
|
# Install OpenCLIP
|
|
print("📥 Installing OpenCLIP...")
|
|
openclip_cmd = [venv_python, '-m', 'pip', 'install', 'open-clip-torch']
|
|
result = subprocess.run(openclip_cmd, capture_output=True, text=True, timeout=300)
|
|
if result.returncode != 0:
|
|
print(f"❌ OpenCLIP installation failed: {result.stderr}")
|
|
return False
|
|
|
|
print("✅ OpenCLIP installed")
|
|
|
|
# Install additional dependencies
|
|
print("📥 Installing additional dependencies...")
|
|
deps_cmd = [venv_python, '-m', 'pip', 'install', 'Pillow', 'requests']
|
|
result = subprocess.run(deps_cmd, capture_output=True, text=True, timeout=300)
|
|
if result.returncode != 0:
|
|
print(f"❌ Dependencies installation failed: {result.stderr}")
|
|
return False
|
|
|
|
print("✅ Additional dependencies installed")
|
|
|
|
# Test the environment
|
|
print("🧪 Testing OpenCLIP GPU environment...")
|
|
test_script = f"""
|
|
import sys
|
|
import torch
|
|
import open_clip
|
|
|
|
print("=== OPENCLIP GPU ENVIRONMENT TEST ===")
|
|
print(f"Python: {{sys.version}}")
|
|
print(f"PyTorch version: {{torch.__version__}}")
|
|
print(f"CUDA available: {{torch.cuda.is_available()}}")
|
|
print(f"CUDA version: {{torch.version.cuda}}")
|
|
print(f"GPU count: {{torch.cuda.device_count()}}")
|
|
|
|
if torch.cuda.is_available():
|
|
for i in range(torch.cuda.device_count()):
|
|
print(f"GPU {{i}}: {{torch.cuda.get_device_name(i)}}")
|
|
|
|
# Test OpenCLIP model loading
|
|
try:
|
|
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
|
|
print("✅ OpenCLIP model loaded successfully")
|
|
|
|
if torch.cuda.is_available():
|
|
model = model.cuda()
|
|
print("✅ OpenCLIP model moved to GPU")
|
|
else:
|
|
print("⚠️ OpenCLIP model running on CPU")
|
|
|
|
except Exception as e:
|
|
print(f"❌ OpenCLIP test failed: {{e}}")
|
|
sys.exit(1)
|
|
"""
|
|
|
|
with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
|
|
f.write(test_script)
|
|
test_file = f.name
|
|
|
|
result = subprocess.run([venv_python, test_file], capture_output=True, text=True, timeout=60)
|
|
os.unlink(test_file)
|
|
|
|
if result.returncode == 0:
|
|
print("🎉 OPENCLIP GPU ENVIRONMENT SETUP SUCCESSFUL!")
|
|
print(result.stdout)
|
|
return True
|
|
else:
|
|
print("❌ OpenCLIP environment test failed")
|
|
print(result.stderr)
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"❌ Environment creation failed: {e}")
|
|
return False
|
|
|
|
def update_image_classifier_for_gpu():
|
|
"""Update the image classifier to use the new GPU environment"""
|
|
print("\n🔄 UPDATING IMAGE CLASSIFIER FOR GPU")
|
|
print("=" * 40)
|
|
|
|
# Update the isolated_image_classifier.py to use the new GPU environment
|
|
new_classifier_code = '''
|
|
"""
|
|
Isolated Image Classifier with GPU Support
|
|
Runs in separate virtual environment with GPU-enabled PyTorch
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import tempfile
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class IsolatedImageClassifier:
|
|
"""Image classifier with GPU support in isolated environment"""
|
|
|
|
def __init__(self):
|
|
self.available = False
|
|
self._check_availability()
|
|
|
|
def _check_availability(self):
|
|
"""Check if OpenCLIP GPU environment is available"""
|
|
try:
|
|
# Use the GPU virtual environment Python executable
|
|
venv_python = "openclip_gpu_env\\\\Scripts\\\\python.exe" if os.name == 'nt' else "openclip_gpu_env/bin/python"
|
|
|
|
if not os.path.exists(venv_python):
|
|
raise RuntimeError(f"GPU virtual environment not found: {venv_python}")
|
|
|
|
# Test OpenCLIP and GPU availability
|
|
test_script = """
|
|
import torch
|
|
import open_clip
|
|
|
|
print(f"GPU_AVAILABLE:{torch.cuda.is_available()}")
|
|
if torch.cuda.is_available():
|
|
print(f"GPU_COUNT:{torch.cuda.device_count()}")
|
|
print(f"GPU_NAME:{torch.cuda.get_device_name(0)}")
|
|
|
|
# Test model loading
|
|
model, _, _ = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
|
|
print("MODEL_LOADED:SUCCESS")
|
|
"""
|
|
result = subprocess.run([venv_python, '-c', test_script],
|
|
capture_output=True, text=True, timeout=30)
|
|
|
|
if "MODEL_LOADED:SUCCESS" in result.stdout:
|
|
self.available = True
|
|
# Parse GPU info
|
|
gpu_available = "GPU_AVAILABLE:True" in result.stdout
|
|
if gpu_available:
|
|
logger.info("✅ OpenCLIP GPU environment initialized with GPU support")
|
|
else:
|
|
logger.warning("⚠️ OpenCLIP GPU environment initialized but running on CPU")
|
|
else:
|
|
error_msg = result.stderr or result.stdout or "Unknown error"
|
|
raise RuntimeError(f"OpenCLIP GPU check failed: {error_msg}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"OpenCLIP GPU availability check failed: {e}")
|
|
raise RuntimeError(f"OpenCLIP GPU availability check failed: {e}")
|
|
|
|
def classify_image(self, image_path: str, top_k: int = 5) -> List[Dict[str, Any]]:
|
|
"""
|
|
Classify image using isolated GPU environment
|
|
|
|
Args:
|
|
image_path: Path to image file
|
|
top_k: Number of top predictions to return
|
|
|
|
Returns:
|
|
List of classification results with confidence scores
|
|
"""
|
|
if not self.available:
|
|
return [{"label": "classification_unavailable", "confidence": 0.0}]
|
|
|
|
try:
|
|
venv_python = "openclip_gpu_env\\\\Scripts\\\\python.exe" if os.name == 'nt' else "openclip_gpu_env/bin/python"
|
|
|
|
classification_script = f"""
|
|
import open_clip
|
|
import torch
|
|
from PIL import Image
|
|
|
|
try:
|
|
# Load model
|
|
model, _, processor = open_clip.create_model_and_transforms(
|
|
model_name="ViT-B-32",
|
|
pretrained="laion2b_s34b_b79k"
|
|
)
|
|
|
|
# Move to GPU if available
|
|
if torch.cuda.is_available():
|
|
model = model.cuda()
|
|
print("USING_GPU:YES")
|
|
else:
|
|
print("USING_GPU:NO")
|
|
|
|
# Load and process image
|
|
image = Image.open(r"{image_path}").convert("RGB")
|
|
image_tensor = processor(image).unsqueeze(0)
|
|
|
|
if torch.cuda.is_available():
|
|
image_tensor = image_tensor.cuda()
|
|
|
|
# Get predictions
|
|
with torch.no_grad():
|
|
image_features = model.encode_image(image_tensor)
|
|
image_features /= image_features.norm(dim=-1, keepdim=True)
|
|
|
|
# Common labels for document images
|
|
text_labels = [
|
|
"a photo of a bee", "a photo of a flower", "a photo of a person",
|
|
"a photo of a document", "a photo of a chart", "a photo of a diagram",
|
|
"a photo of a table", "a photo of a graph", "a photo of a logo",
|
|
"a photo of a signature", "a photo of a stamp", "a photo of a barcode",
|
|
"a photo of a QR code", "a photo of a screenshot", "a photo of a landscape",
|
|
"a photo of an animal", "a photo of a building", "a photo of a vehicle",
|
|
"a photo of food", "a photo of clothing", "a photo of electronics",
|
|
"a photo of furniture", "a photo of nature", "a photo of art",
|
|
"a photo of text", "a photo of numbers", "a photo of symbols"
|
|
]
|
|
|
|
# Encode text labels
|
|
text_tokens = open_clip.tokenize(text_labels)
|
|
if torch.cuda.is_available():
|
|
text_tokens = text_tokens.cuda()
|
|
|
|
text_features = model.encode_text(text_tokens)
|
|
text_features /= text_features.norm(dim=-1, keepdim=True)
|
|
|
|
# Calculate similarity
|
|
similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
|
|
values, indices = similarity[0].topk({top_k})
|
|
|
|
for value, index in zip(values, indices):
|
|
label = text_labels[index]
|
|
confidence = float(value)
|
|
print(f"CLASSIFICATION_RESULT:{{label}}|{{confidence}}")
|
|
|
|
except Exception as e:
|
|
print(f"CLASSIFICATION_ERROR:{{e}}")
|
|
"""
|
|
result = subprocess.run([venv_python, '-c', classification_script],
|
|
capture_output=True, text=True, timeout=30)
|
|
|
|
results = []
|
|
gpu_used = False
|
|
|
|
for line in result.stdout.split('\\n'):
|
|
if line.startswith("USING_GPU:YES"):
|
|
gpu_used = True
|
|
if line.startswith("CLASSIFICATION_RESULT:"):
|
|
parts = line.replace("CLASSIFICATION_RESULT:", "").split("|")
|
|
if len(parts) == 2:
|
|
try:
|
|
label = parts[0]
|
|
confidence = float(parts[1])
|
|
results.append({{"label": label, "confidence": confidence, "gpu_used": gpu_used}})
|
|
except:
|
|
continue
|
|
|
|
if results:
|
|
if gpu_used:
|
|
logger.info(f"✅ Classification completed with GPU on {image_path}")
|
|
else:
|
|
logger.warning(f"⚠️ Classification completed on CPU for {image_path}")
|
|
return results
|
|
else:
|
|
logger.error("No results from classification")
|
|
return [{{"label": "classification_failed", "confidence": 0.0, "gpu_used": False}}]
|
|
|
|
except Exception as e:
|
|
logger.error(f"Classification failed: {e}")
|
|
return [{{"label": "classification_error", "confidence": 0.0, "gpu_used": False}}]
|
|
|
|
|
|
# Singleton instance
|
|
_classifier_instance = None
|
|
|
|
def get_isolated_classifier() -> IsolatedImageClassifier:
|
|
"""Get singleton isolated image classifier instance"""
|
|
global _classifier_instance
|
|
if _classifier_instance is None:
|
|
_classifier_instance = IsolatedImageClassifier()
|
|
return _classifier_instance
|
|
|
|
|
|
def test_isolated_classifier():
|
|
"""Test function for isolated image classifier"""
|
|
classifier = get_isolated_classifier()
|
|
|
|
if classifier.available:
|
|
print("✅ Isolated image classifier is available")
|
|
# Test with a sample image if available
|
|
test_images = ["test_bee_image.png", "sample_image.jpg"]
|
|
for test_image in test_images:
|
|
if os.path.exists(test_image):
|
|
results = classifier.classify_image(test_image)
|
|
print(f"Classification results for {test_image}:")
|
|
for result in results:
|
|
print(f" {result['label']}: {result['confidence']:.4f} (GPU: {result.get('gpu_used', False)})")
|
|
break
|
|
else:
|
|
print("⚠️ No test images found for classification")
|
|
else:
|
|
print("❌ Isolated image classifier is not available")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_isolated_classifier()
|
|
'''
|
|
|
|
with open('isolated_image_classifier_gpu.py', 'w') as f:
|
|
f.write(new_classifier_code)
|
|
|
|
print("✅ Updated image classifier created: isolated_image_classifier_gpu.py")
|
|
return True
|
|
|
|
def main():
|
|
"""Main function to fix OpenCLIP GPU dependencies"""
|
|
print("🚀 FIXING OPENCLIP GPU DEPENDENCIES")
|
|
print("=" * 50)
|
|
|
|
# Step 1: Check system GPU
|
|
if not check_system_gpu():
|
|
print("❌ System GPU check failed. Cannot proceed with GPU setup.")
|
|
return False
|
|
|
|
# Step 2: Create OpenCLIP GPU environment
|
|
if not create_openclip_gpu_environment():
|
|
print("❌ Failed to create OpenCLIP GPU environment")
|
|
return False
|
|
|
|
# Step 3: Update image classifier for GPU
|
|
if not update_image_classifier_for_gpu():
|
|
print("❌ Failed to update image classifier for GPU")
|
|
return False
|
|
|
|
print("\n🎉 OPENCLIP GPU DEPENDENCIES FIXED SUCCESSFULLY!")
|
|
print("=" * 50)
|
|
print("✅ New virtual environment: openclip_gpu_env")
|
|
print("✅ GPU-enabled PyTorch installed")
|
|
print("✅ OpenCLIP with GPU support configured")
|
|
print("✅ Updated image classifier: isolated_image_classifier_gpu.py")
|
|
print("\n📝 Next steps:")
|
|
print(" 1. Update your document processor to use isolated_image_classifier_gpu.py")
|
|
print(" 2. Run test.py to verify GPU usage for both PaddleOCR and OpenCLIP")
|
|
|
|
return True
|
|
|
|
if __name__ == "__main__":
|
|
success = main()
|
|
sys.exit(0 if success else 1) |