File size: 2,434 Bytes
d9b7396 4e5ac0b d9b7396 a1da087 d9b7396 7cb4199 550eca9 d9b7396 a1da087 d9b7396 a1da087 d9b7396 |
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 |
---
license: mit
language:
- en
base_model:
- microsoft/Florence-2-large
datasets:
- diffusers/ShotDEAD-v0
---
# Shot Categorizer 🎬
<div align="center">
<img src="assets/header.jpg"/>
</div>
Shot categorization model finetuned from the [`microsoft/Florence-2-large`](https://huggingface.co/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](https://huggingface.co/datasets/diffusers/ShotDEAD-v0)
dataset for conducting training.
Training code is available [here](https://github.com/huggingface/movie-shot-categorizer).
## Inference
```py
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:
```bash
{'<COLOR>': 'Cool, Saturated, Cyan, Blue'}
{'<LIGHTING>': 'Soft light, Low contrast'}
{'<LIGHTING_TYPE>': 'Daylight, Sunny'}
{'<COMPOSITION>': 'Left heavy'}
``` |