File size: 7,406 Bytes
138d2a7 | 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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | # predict.py
import argparse
import os
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
from transformers import CLIPProcessor, CLIPModel
from peft import LoraConfig, get_peft_model
from typing import Dict, Any
# =========================================================================
# Re-define the Model Class (must match your training script exactly)
# =========================================================================
# Global constants from your training script
NUM_CLASSES = 3
CLIP_MODEL_NAME = "openai/clip-vit-large-patch14"
LORA_R = 32
LORA_ALPHA = 64
LORA_DROPOUT = 0.1
LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"]
FEATURE_COMBINATION_STRATEGY = "mean"
LEARNING_RATE = 2e-4
CLASS_NAMES = ['10k-100k', '100k-1M', '1M+']
class CLIPForViewsClassification(pl.LightningModule):
def __init__(self, num_classes, clip_model_name, lora_r, lora_alpha, lora_dropout, lora_target_modules, learning_rate, feature_combination_strategy, num_training_steps_total):
super().__init__()
self.save_hyperparameters()
self.num_classes = num_classes
self.learning_rate = learning_rate
self.feature_combination_strategy = feature_combination_strategy
self.num_training_steps_total = num_training_steps_total
self.clip_model = CLIPModel.from_pretrained(clip_model_name)
lora_config = LoraConfig(r=lora_r, lora_alpha=lora_alpha, target_modules=lora_target_modules, lora_dropout=lora_dropout, bias="none")
self.clip_model.vision_model = get_peft_model(self.clip_model.vision_model, lora_config)
self.clip_model.text_model = get_peft_model(self.clip_model.text_model, lora_config)
embedding_dim = self.clip_model.config.projection_dim
classifier_input_dim = embedding_dim * 2 if self.feature_combination_strategy == "concat" else embedding_dim
self.classifier = nn.Sequential(
nn.Linear(classifier_input_dim, 512),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(512, self.num_classes)
)
self.processor = CLIPProcessor.from_pretrained(clip_model_name)
def forward(self, pixel_values, input_ids, attention_mask):
image_embeds = self.clip_model.get_image_features(pixel_values=pixel_values)
text_embeds = self.clip_model.get_text_features(input_ids=input_ids, attention_mask=attention_mask)
# --- THE CORRECTED FIX ---
# If the model returns a dataclass instead of a raw tensor,
# the final projected embeddings are ALREADY inside the pooler_output.
# We just extract them safely without projecting them a second time.
if not isinstance(image_embeds, torch.Tensor):
image_embeds = getattr(image_embeds, "pooler_output", image_embeds)
if not isinstance(text_embeds, torch.Tensor):
text_embeds = getattr(text_embeds, "pooler_output", text_embeds)
# -----------------------------------
# Now both are guaranteed to be standard PyTorch tensors
image_embeds = F.normalize(image_embeds, p=2, dim=-1)
text_embeds = F.normalize(text_embeds, p=2, dim=-1)
if self.feature_combination_strategy == "mean":
combined_features = (image_embeds + text_embeds) / 2
else:
combined_features = torch.cat((image_embeds, text_embeds), dim=-1)
logits = self.classifier(combined_features)
return logits
def predict_step(self, image_path: str, title: str) -> Dict[str, Any]:
image = Image.open(image_path).convert("RGB")
inputs = self.processor(text=[title], images=[image], return_tensors="pt", padding="max_length", truncation=True)
device = self.device
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
logits = self(inputs['pixel_values'], inputs['input_ids'], inputs['attention_mask'])
probabilities = F.softmax(logits, dim=1).squeeze(0)
predicted_class_idx = torch.argmax(probabilities).item()
predicted_class = CLASS_NAMES[predicted_class_idx]
return {
"image_path": image_path,
"title": title,
"predicted_class": predicted_class,
"probabilities": probabilities.cpu().numpy()
}
# =========================================================================
# Main Script Logic
# =========================================================================
def main():
parser = argparse.ArgumentParser(description="Use a fine-tuned CLIP model to predict YouTube view categories.")
parser.add_argument("--model_path", type=str, required=True,
help="Path to the trained model checkpoint (.ckpt).")
parser.add_argument("--input_path", type=str, required=True,
help="Path to a single image file or a directory of images to predict on.")
parser.add_argument("--title", type=str, default="A YouTube video",
help="The video title to be used for prediction. Default is 'A YouTube video'.")
args = parser.parse_args()
# Load the trained model
try:
model = CLIPForViewsClassification.load_from_checkpoint(
args.model_path,
num_classes=NUM_CLASSES,
clip_model_name=CLIP_MODEL_NAME,
lora_r=LORA_R,
lora_alpha=LORA_ALPHA,
lora_dropout=LORA_DROPOUT,
lora_target_modules=LORA_TARGET_MODULES,
learning_rate=LEARNING_RATE,
feature_combination_strategy=FEATURE_COMBINATION_STRATEGY,
num_training_steps_total=1000 # Placeholder, not used for inference
)
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f"Model loaded successfully from {args.model_path} and moved to {device}.")
except Exception as e:
print(f"Error loading model: {e}")
return
# Check if input path is a file or directory
if os.path.isfile(args.input_path):
image_paths = [args.input_path]
elif os.path.isdir(args.input_path):
image_paths = [
os.path.join(args.input_path, f)
for f in os.listdir(args.input_path)
if f.lower().endswith(('.png', '.jpg', '.jpeg'))
]
if not image_paths:
print(f"No valid images found in directory: {args.input_path}")
return
else:
print(f"Error: Invalid input path. Must be a valid file or directory.")
return
# Make predictions
print("\n--- Predictions ---")
for img_path in image_paths:
try:
prediction_result = model.predict_step(img_path, args.title)
print(f"\nImage: {prediction_result['image_path']}")
print(f"Predicted Class: {prediction_result['predicted_class']}")
print(f"Probabilities: {prediction_result['probabilities']}")
except Exception as e:
print(f"Error processing image {img_path}: {e}")
if __name__ == "__main__":
main() |