378 lines
18 KiB
Python
378 lines
18 KiB
Python
"""
|
|
Fast Image Classifier using OpenCLIP with GPU acceleration
|
|
Optimized for batch processing and persistent GPU memory usage
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import logging
|
|
import subprocess
|
|
import json
|
|
import tempfile
|
|
from typing import List, Dict, Any, Optional
|
|
from pathlib import Path
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class FastImageClassifier:
|
|
"""Fast image classifier using persistent OpenCLIP process"""
|
|
|
|
def __init__(self):
|
|
self.available = False
|
|
self._process = None
|
|
self._temp_dir = None
|
|
self._initialize_classifier()
|
|
|
|
def _initialize_classifier(self):
|
|
"""Initialize the persistent classifier process"""
|
|
try:
|
|
# Create temporary directory for communication
|
|
self._temp_dir = tempfile.mkdtemp(prefix="openclip_")
|
|
logger.info(f"Created temporary directory for OpenCLIP: {self._temp_dir}")
|
|
|
|
# Start persistent classifier process
|
|
classifier_script = """
|
|
import sys
|
|
import os
|
|
import json
|
|
import torch
|
|
import open_clip
|
|
from PIL import Image
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
# Initialize OpenCLIP model
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
|
|
model = model.to(device)
|
|
tokenizer = open_clip.get_tokenizer('ViT-B-32')
|
|
|
|
# Define labels for classification
|
|
labels = [
|
|
"a photo of a bee", "a photo of a flower", "a photo of a car",
|
|
"a photo of a person", "a photo of a dog", "a photo of a cat",
|
|
"a photo of a building", "a photo of a tree", "a photo of a bird",
|
|
"a photo of a computer", "a photo of a phone", "a photo of a book",
|
|
"a photo of a chair", "a photo of a table", "a photo of a cloud",
|
|
"a photo of the sun", "a photo of the moon", "a photo of a star",
|
|
"a photo of water", "a photo of fire", "a photo of earth",
|
|
"a photo of a mountain", "a photo of a river", "a photo of a lake",
|
|
"a photo of a ocean", "a photo of a desert", "a photo of a forest",
|
|
"a photo of a city", "a photo of a village", "a photo of a road",
|
|
"a photo of a bridge", "a photo of a train", "a photo of a plane",
|
|
"a photo of a boat", "a photo of a bicycle", "a photo of a motorcycle",
|
|
"a photo of a bus", "a photo of a truck", "a photo of a horse",
|
|
"a photo of a cow", "a photo of a sheep", "a photo of a goat",
|
|
"a photo of a pig", "a photo of a chicken", "a photo of a duck",
|
|
"a photo of a fish", "a photo of a shark", "a photo of a whale",
|
|
"a photo of a dolphin", "a photo of a elephant", "a photo of a lion",
|
|
"a photo of a tiger", "a photo of a bear", "a photo of a wolf",
|
|
"a photo of a fox", "a photo of a rabbit", "a photo of a deer",
|
|
"a photo of a monkey", "a photo of a snake", "a photo of a lizard",
|
|
"a photo of a frog", "a photo of a turtle", "a photo of a spider",
|
|
"a photo of a insect", "a photo of a butterfly", "a photo of a dragonfly",
|
|
"a photo of a ant", "a photo of a fly", "a photo of a mosquito",
|
|
"a photo of a grasshopper", "a photo of a cricket", "a photo of a beetle",
|
|
"a photo of a ladybug", "a photo of a cockroach", "a photo of a scorpion",
|
|
"a photo of a centipede", "a photo of a millipede", "a photo of a worm",
|
|
"a photo of a snail", "a photo of a slug", "a photo of a octopus",
|
|
"a photo of a squid", "a photo of a jellyfish", "a photo of a coral",
|
|
"a photo of a seaweed", "a photo of a mushroom", "a photo of a plant",
|
|
"a photo of a leaf", "a photo of a root", "a photo of a seed",
|
|
"a photo of a fruit", "a photo of a vegetable", "a photo of a grain",
|
|
"a photo of a nut", "a photo of a berry", "a photo of a flower",
|
|
"a photo of a rose", "a photo of a tulip", "a photo of a sunflower",
|
|
"a photo of a lily", "a photo of a orchid", "a photo of a daisy",
|
|
"a photo of a dandelion", "a photo of a cactus", "a photo of a palm tree",
|
|
"a photo of a pine tree", "a photo of a oak tree", "a photo of a maple tree",
|
|
"a photo of a birch tree", "a photo of a willow tree", "a photo of a cherry blossom",
|
|
"a photo of a bamboo", "a photo of a grass", "a photo of a moss",
|
|
"a photo of a fern", "a photo of a ivy", "a photo of a vine",
|
|
"a photo of a weed", "a photo of a herb", "a photo of a spice",
|
|
"a photo of a salt", "a photo of a pepper", "a photo of a sugar",
|
|
"a photo of a flour", "a photo of a bread", "a photo of a cheese",
|
|
"a photo of a milk", "a photo of a egg", "a photo of a meat",
|
|
"a photo of a fish", "a photo of a chicken", "a photo of a beef",
|
|
"a photo of a pork", "a photo of a lamb", "a photo of a turkey",
|
|
"a photo of a duck", "a photo of a goose", "a photo of a rabbit",
|
|
"a photo of a deer", "a photo of a bison", "a photo of a buffalo",
|
|
"a photo of a elk", "a photo of a moose", "a photo of a caribou",
|
|
"a photo of a antelope", "a photo of a gazelle", "a photo of a zebra",
|
|
"a photo of a giraffe", "a photo of a hippopotamus", "a photo of a rhinoceros",
|
|
"a photo of a crocodile", "a photo of a alligator", "a photo of a turtle",
|
|
"a photo of a tortoise", "a photo of a lizard", "a photo of a snake",
|
|
"a photo of a frog", "a photo of a toad", "a photo of a salamander",
|
|
"a photo of a newt", "a photo of a caecilian", "a photo of a fish",
|
|
"a photo of a shark", "a photo of a ray", "a photo of a eel",
|
|
"a photo of a octopus", "a photo of a squid", "a photo of a cuttlefish",
|
|
"a photo of a nautilus", "a photo of a jellyfish", "a photo of a coral",
|
|
"a photo of a sponge", "a photo of a sea anemone", "a photo of a starfish",
|
|
"a photo of a sea urchin", "a photo of a sea cucumber", "a photo of a crab",
|
|
"a photo of a lobster", "a photo of a shrimp", "a photo of a prawn",
|
|
"a photo of a crayfish", "a photo of a barnacle", "a photo of a mussel",
|
|
"a photo of a clam", "a photo of a oyster", "a photo of a scallop",
|
|
"a photo of a snail", "a photo of a slug", "a photo of a worm",
|
|
"a photo of a leech", "a photo of a centipede", "a photo of a millipede",
|
|
"a photo of a spider", "a photo of a scorpion", "a photo of a tick",
|
|
"a photo of a mite", "a photo of a flea", "a photo of a louse",
|
|
"a photo of a bedbug", "a photo of a cockroach", "a photo of a termite",
|
|
"a photo of a ant", "a photo of a bee", "a photo of a wasp",
|
|
"a photo of a hornet", "a photo of a yellowjacket", "a photo of a fly",
|
|
"a photo of a mosquito", "a photo of a gnat", "a photo of a midge",
|
|
"a photo of a butterfly", "a photo of a moth", "a photo of a dragonfly",
|
|
"a photo of a damselfly", "a photo of a grasshopper", "a photo of a cricket",
|
|
"a photo of a katydid", "a photo of a cicada", "a photo of a beetle",
|
|
"a photo of a ladybug", "a photo of a firefly", "a photo of a weevil",
|
|
"a photo of a stag beetle", "a photo of a rhinoceros beetle", "a photo of a dung beetle",
|
|
"a photo of a ground beetle", "a photo of a water beetle", "a photo of a leaf beetle",
|
|
"a photo of a longhorn beetle", "a photo of a jewel beetle", "a photo of a darkling beetle",
|
|
"a photo of a click beetle", "a photo of a glowworm", "a photo of a earwig",
|
|
"a photo of a silverfish", "a photo of a springtail", "a photo of a proturan",
|
|
"a photo of a dipluran", "a photo of a collembolan", "a photo of a thrips",
|
|
"a photo of a aphid", "a photo of a scale insect", "a photo of a mealybug",
|
|
"a photo of a whitefly", "a photo of a psyllid", "a photo of a leafhopper",
|
|
"a photo of a treehopper", "a photo of a froghopper", "a photo of a planthopper",
|
|
"a photo of a cicada", "a photo of a spittlebug", "a photo of a lanternfly",
|
|
"a photo of a stink bug", "a photo of a shield bug", "a photo of a assassin bug",
|
|
"a photo of a bed bug", "a photo of a water strider", "a photo of a backswimmer",
|
|
"a photo of a water boatman", "a photo of a giant water bug", "a photo of a kissing bug",
|
|
"a photo of a wheel bug", "a photo of a leaf-footed bug", "a photo of a squash bug",
|
|
"a photo of a boxelder bug", "a photo of a milkweed bug", "a photo of a seed bug",
|
|
"a photo of a lace bug", "a photo of a plant bug", "a photo of a mirid bug",
|
|
"a photo of a capsid bug", "a photo of a damsel bug", "a photo of a minute pirate bug",
|
|
"a photo of a flower bug", "a photo of a anthocorid bug", "a photo of a reduviid bug",
|
|
"a photo of a tingid bug", "a photo of a coreid bug", "a photo of a lygaeid bug",
|
|
"a photo of a pyrrhocorid bug", "a photo of a alydid bug", "a photo of a rhopalid bug",
|
|
"a photo of a pentatomid bug", "a photo of a scutellerid bug", "a photo of a cydnid bug",
|
|
"a photo of a thynnid bug", "a photo of a plataspid bug", "a photo of a acanthosomatid bug",
|
|
"a photo of a urostylid bug", "a photo of a aradid bug", "a photo of a termite",
|
|
"a photo of a cockroach", "a photo of a mantis", "a photo of a stick insect",
|
|
"a photo of a leaf insect", "a photo of a grasshopper", "a photo of a cricket",
|
|
"a photo of a katydid", "a photo of a cicada", "a photo of a planthopper",
|
|
"a photo of a treehopper", "a photo of a froghopper", "a photo of a leafhopper",
|
|
"a photo of a spittlebug", "a photo of a lanternfly", "a photo of a aphid",
|
|
"a photo of a scale insect", "a photo of a mealybug", "a photo of a whitefly",
|
|
"a photo of a psyllid", "a photo of a thrips", "a photo of a lacewing",
|
|
"a photo of a antlion", "a photo of a dobsonfly", "a photo of a fishfly",
|
|
"a photo of a alderfly", "a photo of a snakefly", "a photo of a scorpionfly",
|
|
"a photo of a hangingfly", "a photo of a caddisfly", "a photo of a butterfly",
|
|
"a photo of a moth", "a photo of a skipper", "a photo of a sawfly",
|
|
"a photo of a horntail", "a photo of a wood wasp", "a photo of a ichneumon wasp",
|
|
"a photo of a braconid wasp", "a photo of a chalcid wasp", "a photo of a gall wasp",
|
|
"a photo of a fig wasp", "a photo of a pollen wasp", "a photo of a velvet ant",
|
|
"a photo of a spider wasp", "a photo of a potter wasp", "a photo of a paper wasp",
|
|
"a photo of a yellowjacket", "a photo of a hornet", "a photo of a honey bee",
|
|
"a photo of a bumblebee", "a photo of a carpenter bee", "a photo of a leafcutter bee",
|
|
"a photo of a mason bee", "a photo of a sweat bee", "a photo of a mining bee",
|
|
"a photo of a plasterer bee", "a photo of a cuckoo bee", "a photo of a orchid bee",
|
|
"a photo of a stingless bee", "a photo of a ant", "a photo of a termite",
|
|
"a photo of a wasp", "a photo of a bee", "a photo of a hornet",
|
|
"a photo of a yellowjacket", "a photo of a sawfly", "a photo of a horntail",
|
|
"a photo of a wood wasp", "a photo of a ichneumon wasp", "a photo of a braconid wasp",
|
|
"a photo of a chalcid wasp", "a photo of a gall wasp", "a photo of a fig wasp",
|
|
"a photo of a pollen wasp", "a photo of a velvet ant", "a photo of a spider wasp",
|
|
"a photo of a potter wasp", "a photo of a paper wasp", "a photo of a yellowjacket",
|
|
"a photo of a hornet", "a photo of a honey bee", "a photo of a bumblebee",
|
|
"a photo of a carpenter bee", "a photo of a leafcutter bee", "a photo of a mason bee",
|
|
"a photo of a sweat bee", "a photo of a mining bee", "a photo of a plasterer bee",
|
|
"a photo of a cuckoo bee", "a photo of a orchid bee", "a photo of a stingless bee",
|
|
"a photo of a ant", "a photo of a termite", "a photo of a wasp",
|
|
"a photo of a bee", "a photo of a hornet", "a photo of a yellowjacket",
|
|
"a photo of a sawfly", "a photo of a horntail", "a photo of a wood wasp",
|
|
"a photo of a ichneumon wasp", "a photo of a braconid wasp", "a photo of a chalcid wasp",
|
|
"a photo of a gall wasp", "a photo of a fig wasp", "a photo of a pollen wasp",
|
|
"a photo of a velvet ant", "a photo of a spider wasp", "a photo of a potter wasp",
|
|
"a photo of a paper wasp", "a photo of a yellowjacket", "a photo of a hornet",
|
|
"a photo of a honey bee", "a photo of a bumblebee", "a photo of a carpenter bee",
|
|
"a photo of a leafcutter bee", "a photo of a mason bee", "a photo of a sweat bee",
|
|
"a photo of a mining bee", "a photo of a plasterer bee", "a photo of a cuckoo bee",
|
|
"a photo of a orchid bee", "a photo of a stingless bee", "a photo of a ant",
|
|
"a photo of a termite", "a photo of a wasp", "a photo of a bee"
|
|
]
|
|
|
|
text = tokenizer(labels).to(device)
|
|
|
|
def classify_image(image_path, top_k=3):
|
|
try:
|
|
# Load and preprocess image
|
|
image = Image.open(image_path).convert('RGB')
|
|
image_input = preprocess(image).unsqueeze(0).to(device)
|
|
|
|
# Extract features
|
|
with torch.no_grad():
|
|
image_features = model.encode_image(image_input)
|
|
text_features = model.encode_text(text)
|
|
|
|
# Normalize features
|
|
image_features = image_features / image_features.norm(dim=-1, keepdim=True)
|
|
text_features = 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)
|
|
|
|
results = []
|
|
for value, index in zip(values, indices):
|
|
results.append({
|
|
"label": labels[index],
|
|
"confidence": float(value)
|
|
})
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
return [{"label": f"Error: {str(e)}", "confidence": 0.0}]
|
|
|
|
# Main loop for processing requests
|
|
while True:
|
|
try:
|
|
line = input().strip()
|
|
if not line:
|
|
continue
|
|
|
|
request = json.loads(line)
|
|
action = request.get("action")
|
|
|
|
if action == "classify":
|
|
image_path = request["image_path"]
|
|
top_k = request.get("top_k", 3)
|
|
results = classify_image(image_path, top_k)
|
|
response = {"success": True, "results": results}
|
|
|
|
elif action == "ping":
|
|
response = {"success": True, "message": "pong"}
|
|
|
|
elif action == "exit":
|
|
break
|
|
|
|
else:
|
|
response = {"success": False, "error": f"Unknown action: {action}"}
|
|
|
|
print(json.dumps(response))
|
|
sys.stdout.flush()
|
|
|
|
except Exception as e:
|
|
error_response = {"success": False, "error": str(e)}
|
|
print(json.dumps(error_response))
|
|
sys.stdout.flush()
|
|
"""
|
|
|
|
# Write the classifier script
|
|
script_path = os.path.join(self._temp_dir, "classifier.py")
|
|
with open(script_path, 'w') as f:
|
|
f.write(classifier_script)
|
|
|
|
# Start the persistent process
|
|
env = os.environ.copy()
|
|
# Use the openclip_gpu_env virtual environment
|
|
python_executable = r"C:\aaWORK\railseek6\openclip_gpu_env\Scripts\python.exe"
|
|
|
|
self._process = subprocess.Popen(
|
|
[python_executable, script_path],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
bufsize=1,
|
|
universal_newlines=True,
|
|
env=env
|
|
)
|
|
|
|
# Test the process
|
|
test_request = {"action": "ping"}
|
|
self._process.stdin.write(json.dumps(test_request) + '\n')
|
|
self._process.stdin.flush()
|
|
response_line = self._process.stdout.readline()
|
|
|
|
try:
|
|
response = json.loads(response_line)
|
|
if response.get("success"):
|
|
self.available = True
|
|
logger.info("Fast image classifier initialized successfully with GPU")
|
|
else:
|
|
logger.error(f"Classifier ping failed: {response.get('error')}")
|
|
except json.JSONDecodeError:
|
|
logger.error("Failed to parse classifier response")
|
|
self.available = False
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize fast image classifier: {e}")
|
|
self.available = False
|
|
|
|
def classify_image(self, image_path: str, top_k: int = 3) -> List[Dict[str, Any]]:
|
|
"""Classify an image and return top_k results"""
|
|
if not self.available or not self._process:
|
|
return [{"label": "Classifier not available", "confidence": 0.0}]
|
|
|
|
try:
|
|
request = {
|
|
"action": "classify",
|
|
"image_path": image_path,
|
|
"top_k": top_k
|
|
}
|
|
|
|
self._process.stdin.write(json.dumps(request) + '\n')
|
|
self._process.stdin.flush()
|
|
response_line = self._process.stdout.readline()
|
|
|
|
response = json.loads(response_line)
|
|
if response.get("success"):
|
|
return response["results"]
|
|
else:
|
|
logger.error(f"Classification failed: {response.get('error')}")
|
|
return [{"label": f"Error: {response.get('error')}", "confidence": 0.0}]
|
|
|
|
except Exception as e:
|
|
logger.error(f"Classification request failed: {e}")
|
|
return [{"label": f"Request error: {str(e)}", "confidence": 0.0}]
|
|
|
|
def close(self):
|
|
"""Close the classifier process"""
|
|
if self._process:
|
|
try:
|
|
# Send exit command
|
|
exit_request = {"action": "exit"}
|
|
self._process.stdin.write(json.dumps(exit_request) + '\n')
|
|
self._process.stdin.flush()
|
|
self._process.wait(timeout=5)
|
|
except:
|
|
self._process.terminate()
|
|
finally:
|
|
self._process = None
|
|
|
|
# Clean up temporary directory
|
|
if self._temp_dir and os.path.exists(self._temp_dir):
|
|
import shutil
|
|
try:
|
|
shutil.rmtree(self._temp_dir)
|
|
except:
|
|
pass
|
|
|
|
def __del__(self):
|
|
"""Destructor to ensure cleanup"""
|
|
self.close()
|
|
|
|
# Singleton instance
|
|
_classifier_instance = None
|
|
|
|
def get_image_classifier() -> FastImageClassifier:
|
|
"""Get singleton classifier instance"""
|
|
global _classifier_instance
|
|
if _classifier_instance is None:
|
|
_classifier_instance = FastImageClassifier()
|
|
return _classifier_instance
|
|
|
|
if __name__ == "__main__":
|
|
# Test the classifier
|
|
classifier = get_image_classifier()
|
|
if classifier.available:
|
|
test_image = "test_image.jpg" # Replace with actual test image
|
|
if os.path.exists(test_image):
|
|
results = classifier.classify_image(test_image)
|
|
print("Classification results:", results)
|
|
else:
|
|
print("Test image not found")
|
|
else:
|
|
print("Classifier not available") |