diffusers / shot-categorizer-v0

Shot Categorizer 🎬

Shot categorization model finetuned from the microsoft/Florence-2-large model. This model can be used to obtain metadata information about shots which can further be used to curate datasets of different kinds.

Training configuration:

• Batch size: 16
• Gradient accumulation steps: 4
• Learning rate: 1e-6
• Epochs: 20
• Max grad norm: 1.0
• Hardware: 8xH100s

Training was conducted using FP16 mixed-precision and DeepSpeed Zero2 scheme. The vision tower of the model was kept frozen during the training. We used the diffusers/ShotDEAD-v0 dataset for conducting training.

Training code is available here .

Inference

from transformers import AutoModelForCausalLM, AutoProcessor
import torch
from PIL import Image
import requests


folder_path = "diffusers/shot-categorizer-v0"
model = (
    AutoModelForCausalLM.from_pretrained(folder_path, torch_dtype=torch.float16, trust_remote_code=True)
    .to("cuda")
    .eval()
)
processor = AutoProcessor.from_pretrained(folder_path, trust_remote_code=True)

prompts = ["<COLOR>", "<LIGHTING>", "<LIGHTING_TYPE>", "<COMPOSITION>"]
img_path = "./assets/image_3.jpg"
image = Image.open(img_path).convert("RGB")

with torch.no_grad() and torch.inference_mode():
    for prompt in prompts:
        inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda", torch.float16)
        generated_ids = model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
        parsed_answer = processor.post_process_generation(
            generated_text, task=prompt, image_size=(image.width, image.height)
        )
        print(parsed_answer)

Should print:

{'<COLOR>': 'Cool, Saturated, Cyan, Blue'}
{'<LIGHTING>': 'Soft light, Low contrast'}
{'<LIGHTING_TYPE>': 'Daylight, Sunny'}
{'<COMPOSITION>': 'Left heavy'}