VLLMClassifier.py

import json
import argparse
from LabelProcessor import LabelProcessor
import torch, os
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageTextToText, AutoProcessor
from qwen_vl_utils import process_vision_info
from MeshRenderer import MeshRenderer

class VLLMClassifier:
    def __init__(self, model_name=None, device="cuda", num_views=12, resolution=1024):
        """Initialize the VLLM-based classifier and supporting components.

        Args:
            model_name (str): Path or identifier for the vision-language model.
            device (str): Device for model execution (e.g. "cuda" or "cpu").
            num_views (int): Number of rendered views to request from the renderer.
            resolution (int): Resolution for render output images.
        """

        self.device = device
        model_name = model_name or os.path.normpath(
            os.path.join(os.path.dirname(__file__), "..", "models", "Qwen3-VL-8B-Instruct")
        )
        self.model = AutoModelForImageTextToText.from_pretrained(model_name, device_map={"": device})
        self.model.eval()
        self.processor = AutoProcessor.from_pretrained(model_name)
        self.mesh_renderer = MeshRenderer(num_views=num_views, resolution=resolution)
        self.classification_prompt = (
            "You are performing object classification. "
            "Based on all provided images of the same object, output exactly ONE label "
            "containing 1 to 3 words. "
            "The label must be the most specific identifiable real-world object category. "
            "If the object represents a well-known fictional character, brand, or copyrighted entity, "
            "output its proper name (e.g., 'Iron Man', 'Spider-Man', 'Batman'). "
            "Do NOT describe shape, color, or geometry. "
            "Do NOT output a sentence. "
            "Do NOT explain your reasoning. "
            "Return only the object name."
        )

    def set_num_views(self, num_views):
        """Update the renderer to use a different number of views.

        Args:
            num_views (int): New number of views to render per mesh.
        """
        self.mesh_renderer = MeshRenderer(num_views=num_views, resolution=self.mesh_renderer.resolution)

    def set_resolution(self, resolution):
        """Update the renderer resolution.

        Args:
            resolution (int): New square pixel resolution for renders.
        """
        self.mesh_renderer = MeshRenderer(num_views=self.mesh_renderer.num_views, resolution=resolution)

    def build_input(self, image_paths):
        """Prepare model inputs (images + prompt) for the vision-language model.

        Args:
            image_paths (list[str]): Paths to images to include as context.

        Returns:
            dict: Tokenized tensors and image tensors suitable for model.generate.
        """
        file_paths = [f"file://{os.path.abspath(path)}" for path in image_paths]
        message = [
            {
                "role": "user",
                "content": (
                    [{"type": "image", "image": path} for path in file_paths]
                    + [{"type": "text", "text": self.classification_prompt}]
                ),
            }
        ]

        text = self.processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
        images, _ = process_vision_info(message, image_patch_size=16)
        inputs = self.processor(text=text, images=images, return_tensors="pt", do_resize=False).to(self.device)
        return inputs


    def classify_one(self, file_path):
        """Render a mesh to images, send them to the model, and decode the label.

        Args:
            file_path (str): Path to the mesh file to classify.

        Returns:
            str: Lowercased predicted label.
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Mesh file {file_path} does not exist.")

        image_paths = self.mesh_renderer.render_views(file_path)
        inputs = self.build_input(image_paths)

        with torch.no_grad():
            outputs = self.model.generate(**inputs, max_new_tokens=128)
    
        # Remove prompt tokens
        generated_ids_trimmed = [
            out_ids[len(in_ids):]
            for in_ids, out_ids in zip(inputs.input_ids, outputs)
        ]

        # Decode
        output_text = self.processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False
        )

        label = output_text[0].strip().lower()

        # Clean up to save GPU space
        del inputs, outputs, generated_ids_trimmed
        torch.cuda.empty_cache()

        print("\nMulti-view classification:")
        print(label)
        return label


    def classify_batch(self, folder_path, limit=None, save_path=None):
        """Classify all meshes in a directory and optionally save results.

        Args:
            folder_path (str): Directory containing mesh files.
            limit (int, optional): Maximum number of files to process.
            save_path (str, optional): Filename to save JSON results under `classifications/`.

        Returns:
            str: Status message including save path when provided.
        """
        if not os.path.isdir(folder_path):
            raise ValueError(f"Folder does not exist: {folder_path}")

        files = []
        for root, _, filenames in os.walk(folder_path):
            files.extend(os.path.join(root, f) for f in filenames)

        files.sort()

        if limit is not None:
            files = files[:limit]
        classifications = []
        label_processor = LabelProcessor()

        for file in files:
            print(f"\nClassifying {file}...")
            predicted_label = self.classify_one(file)
            print(f"Predicted label: {predicted_label}")
            embedding = label_processor.compute_embedding(predicted_label)
            classifications.append({
                'file_path': file,
                'predicted_label': predicted_label,
                'embedding': embedding.tolist(),
            })

        if save_path:
            os.makedirs("classifications", exist_ok=True)
            save_path = os.path.join("classifications", save_path)
            with open(save_path, 'w') as f:
                json.dump(classifications, f)

        return 'Classification completed. Results saved to ' + save_path if save_path else 'Classification completed.'

    def classify(self, file_path=None, folder_path=None, limit=None, save_path=None):
        """Dispatch to single-file or batch classification.

        Args:
            file_path (str, optional): Single mesh file to classify.
            folder_path (str, optional): Directory of mesh files to classify.
            limit (int, optional): Limit for batch processing.
            save_path (str, optional): Path to save batch results.

        Returns:
            Any: Return value from `classify_one` or `classify_batch`.

        Raises:
            ValueError: If neither `file_path` nor `folder_path` is provided.
        """
        if file_path is not None:
            return self.classify_one(file_path)
        elif folder_path is not None:
            return self.classify_batch(folder_path, limit=limit, save_path=save_path)
        else:
            raise ValueError("Either file_path or folder_path must be provided.")

    @classmethod
    def from_cli(cls):
        """Parse command-line arguments and run classification."""
        parser = argparse.ArgumentParser(
            description="Classify 3D meshes using VLLM Classifier",
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog="""
            Examples:
            # Classify a single file
            python VLLMClassifier.py --file path/to/mesh.obj
            
            # Classify all files in a directory
            python VLLMClassifier.py --dir /path/to/meshes --output results.json
            
            # Classify with custom settings
            python VLLMClassifier.py --dir /path/to/meshes --num-views 8 --limit 100 --output results.json
            """
        )
        
        # Input arguments (mutually exclusive)
        input_group = parser.add_mutually_exclusive_group(required=True)
        input_group.add_argument("--file", type=str, help="Path to single mesh file")
        input_group.add_argument("--dir", type=str, help="Path to directory of mesh files")
        
        # Output arguments
        parser.add_argument("--output", type=str, default=None, help="Output file path for batch classifications (saved to classifications/ directory)")
        parser.add_argument("--limit", type=int, default=None, help="Limit number of files to classify (batch mode only)")
        
        # Model arguments
        parser.add_argument("--device", type=str, default="cuda:0", help="Device to use (default: cuda:0)")
        parser.add_argument(
            "--model",
            type=str,
            default=os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "models", "Qwen3-VL-8B-Instruct")),
            help="Path to model",
        )
        parser.add_argument("--num-views", type=int, default=12, help="Number of rendering views (default: 12)")
        parser.add_argument("--resolution", type=int, default=1024, help="Rendering resolution (default: 1024)")
        
        args = parser.parse_args()
        
        # Initialize classifier
        classifier = cls(
            model_name=args.model,
            device=args.device,
            num_views=args.num_views,
            resolution=args.resolution
        )
        
        # Run classification
        if args.file:
            print(f"\nClassifying single file: {args.file}")
            result = classifier.classify_one(args.file)
            print(f"\nFinal classification: {result}")
        else:
            print(f"\nClassifying directory: {args.dir}")
            result = classifier.classify_batch(
                folder_path=args.dir,
                limit=args.limit,
                save_path=args.output
            )
            print(f"\n{result}")
        
        torch.cuda.empty_cache()


if __name__ == "__main__":
    VLLMClassifier.from_cli()