engine/generate_features.py

import asyncio
import os
import logging
from PIL import Image
import torch
from transformers import (
    CLIPProcessor,
    CLIPModel,
    BlipProcessor,
    BlipForConditionalGeneration,
)
from sentence_transformers import SentenceTransformer
import numpy as np
import aiofiles
import json
from abc import ABC, abstractmethod
from typing import Set, Tuple
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

device = "cpu"


@dataclass
class State:
    processed_files: Set[str] = field(default_factory=set)

    def to_dict(self) -> dict:
        return {"processed_files": list(self.processed_files)}

    @staticmethod
    def from_dict(state_dict: dict) -> "State":
        return State(processed_files=set(state_dict.get("processed_files", [])))


class ImageProcessor(ABC):
    @abstractmethod
    def process(self, image: Image.Image) -> np.ndarray:
        pass


class CLIPImageProcessor(ImageProcessor):
    def __init__(self):
        self.model = CLIPModel.from_pretrained(
            "wkcn/TinyCLIP-ViT-8M-16-Text-3M-YFCC15M"
        ).to(device)
        self.processor = CLIPProcessor.from_pretrained(
            "wkcn/TinyCLIP-ViT-8M-16-Text-3M-YFCC15M"
        )
        print("Initialized CLIP model and processor")

    def process(self, image: Image.Image) -> np.ndarray:
        inputs = self.processor(images=image, return_tensors="pt").to(device)
        outputs = self.model.get_image_features(**inputs)
        return outputs.detach().cpu().numpy()


class ImageCaptioningProcessor(ImageProcessor):
    def __init__(self):
        self.image_caption_model = BlipForConditionalGeneration.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        ).to(device)
        self.image_caption_processor = BlipProcessor.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        )
        self.text_embedding_model = SentenceTransformer(
            "all-MiniLM-L6-v2", device=device
        )
        print("Initialized BLIP model and processor")

    def process(self, image: Image.Image) -> np.ndarray:
        inputs = self.image_caption_processor(images=image, return_tensors="pt").to(
            device
        )
        output = self.image_caption_model.generate(**inputs)
        caption = self.image_caption_processor.decode(
            output[0], skip_special_tokens=True
        )
        # embedding dim 384
        return self.text_embedding_model.encode(caption).flatten()


class ImageFeatureExtractor:
    def __init__(
        self,
        clip_processor: CLIPImageProcessor,
        caption_processor: ImageCaptioningProcessor,
        max_queue_size: int = 100,
        checkpoint_file: str = "checkpoint.json",
    ):
        self.clip_processor = clip_processor
        self.caption_processor = caption_processor
        self.image_queue = asyncio.Queue(maxsize=max_queue_size)
        self.processed_images_queue = asyncio.Queue()
        self.checkpoint_file = checkpoint_file
        self.state = self.load_state()
        self.executor = ProcessPoolExecutor()
        self.total_images = 0
        self.processed_count = 0
        print(
            "Initialized ImageFeatureExtractor with checkpoint file:", checkpoint_file
        )

    async def image_loader(self, input_folder: str):
        print(f"Loading images from {input_folder}")
        for filename in os.listdir(input_folder):
            if "resized_" in filename and filename not in self.state.processed_files:
                try:
                    file_path = os.path.join(input_folder, filename)
                    await self.image_queue.put((filename, file_path))
                    self.total_images += 1
                    print(f"Loaded image {filename} into queue")
                except Exception as e:
                    logger.error(f"Error loading image {filename}: {e}")
        await self.image_queue.put(None)  # Sentinel to signal end of images
        print(f"Total images to process: {self.total_images}")

    async def image_processor_worker(self, loop: asyncio.AbstractEventLoop):
        while True:
            item = await self.image_queue.get()
            if item is None:
                await self.image_queue.put(None)  # Propagate sentinel
                break
            filename, file_path = item
            try:
                print(f"Processing image {filename}")
                image = Image.open(file_path)
                clip_embedding, caption_embedding = await asyncio.gather(
                    loop.run_in_executor(
                        self.executor, self.clip_processor.process, image
                    ),
                    loop.run_in_executor(
                        self.executor, self.caption_processor.process, image
                    ),
                )
                await self.processed_images_queue.put(
                    (filename, clip_embedding, caption_embedding)
                )
                print(f"Processed image {filename}")
            except Exception as e:
                logger.error(f"Error processing image {filename}: {e}")
            finally:
                self.image_queue.task_done()

    async def save_processed_images(self, output_folder: str):
        while self.processed_count < self.total_images:
            filename, clip_embedding, caption_embedding = (
                await self.processed_images_queue.get()
            )
            try:
                clip_output_path = os.path.join(
                    output_folder, f"{os.path.splitext(filename)[0]}_clip.npy"
                )
                caption_output_path = os.path.join(
                    output_folder, f"{os.path.splitext(filename)[0]}_caption.npy"
                )

                await asyncio.gather(
                    self.save_embedding(clip_output_path, clip_embedding),
                    self.save_embedding(caption_output_path, caption_embedding),
                )

                self.state.processed_files.add(filename)
                self.save_state()
                self.processed_count += 1
                print(f"Saved processed embeddings for {filename}")
            except Exception as e:
                logger.error(f"Error saving processed image {filename}: {e}")
            finally:
                self.processed_images_queue.task_done()

    async def save_embedding(self, output_path: str, embedding: np.ndarray):
        async with aiofiles.open(output_path, "wb") as f:
            await f.write(embedding.tobytes())

    def load_state(self) -> State:
        try:
            with open(self.checkpoint_file, "r") as f:
                state_dict = json.load(f)
                print("Loaded state from checkpoint")
                return State.from_dict(state_dict)
        except (FileNotFoundError, json.JSONDecodeError):
            print("No checkpoint found, starting with empty state")
            return State()

    def save_state(self):
        with open(self.checkpoint_file, "w") as f:
            json.dump(self.state.to_dict(), f)
            print("Saved state to checkpoint")

    async def run(
        self,
        input_folder: str,
        output_folder: str,
        loop: asyncio.AbstractEventLoop,
        num_workers: int = 2,
    ):
        os.makedirs(output_folder, exist_ok=True)
        print(f"Output folder {output_folder} created")

        tasks = [
            loop.create_task(self.image_loader(input_folder)),
            loop.create_task(self.save_processed_images(output_folder)),
        ]
        tasks.extend(
            [
                loop.create_task(self.image_processor_worker(loop))
                for _ in range(num_workers)
            ]
        )

        await asyncio.gather(*tasks)


class ImageFeatureExtractorFactory:
    @staticmethod
    def create() -> ImageFeatureExtractor:
        print(
            "Creating ImageFeatureExtractor with CLIPImageProcessor and ImageCaptioningProcessor"
        )
        return ImageFeatureExtractor(CLIPImageProcessor(), ImageCaptioningProcessor())


async def main(loop: asyncio.AbstractEventLoop, input_folder: str, output_folder: str):
    print("Starting main function")

    extractor = ImageFeatureExtractorFactory.create()

    try:
        await extractor.run(input_folder, output_folder, loop)
    except Exception as e:
        logger.error(f"An error occurred during execution: {e}")
    finally:
        logger.info("Image processing completed.")


if __name__ == "__main__":
    from pathlib import Path

    PROJECT_ROOT = Path(__file__).resolve().parent.parent
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    print("Event loop created and set")
    input_folder = str(PROJECT_ROOT / "data/images")
    output_folder = str(PROJECT_ROOT / "data/features")
    loop.run_until_complete(main(loop, input_folder, output_folder))
    loop.close()
    print("Event loop closed")