V 3.3.9（2025.05.29）
1. 生图包适配，修复部分任务无法获取图片的问题。 2. 新增聚合推文原创重置功能，直接重置GPT提示词、合并提示词、出图等。 3. 优化加载界面，提升用户体验。 4. 优化主页显示，优化UI界面，提升用户体验。
This commit is contained in:
Binary file not shown.
Binary file not shown.
@@ -1,29 +0,0 @@
|
||||
Collecting accelerate
|
||||
Downloading accelerate-1.0.1-py3-none-any.whl.metadata (19 kB)
|
||||
Requirement already satisfied: numpy<3.0.0,>=1.17 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (2.1.2)
|
||||
Requirement already satisfied: packaging>=20.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (24.1)
|
||||
Collecting psutil (from accelerate)
|
||||
Downloading psutil-6.1.0-cp37-abi3-win_amd64.whl.metadata (23 kB)
|
||||
Requirement already satisfied: pyyaml in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (6.0.2)
|
||||
Requirement already satisfied: torch>=1.10.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (2.5.0)
|
||||
Requirement already satisfied: huggingface-hub>=0.21.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (0.26.1)
|
||||
Requirement already satisfied: safetensors>=0.4.3 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (0.4.5)
|
||||
Requirement already satisfied: filelock in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (3.16.1)
|
||||
Requirement already satisfied: fsspec>=2023.5.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (2024.10.0)
|
||||
Requirement already satisfied: requests in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (2.32.3)
|
||||
Requirement already satisfied: tqdm>=4.42.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (4.66.5)
|
||||
Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (4.12.2)
|
||||
Requirement already satisfied: networkx in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (3.4.2)
|
||||
Requirement already satisfied: jinja2 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (3.1.4)
|
||||
Requirement already satisfied: sympy==1.13.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (1.13.1)
|
||||
Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from sympy==1.13.1->torch>=1.10.0->accelerate) (1.3.0)
|
||||
Requirement already satisfied: colorama in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from tqdm>=4.42.1->huggingface-hub>=0.21.0->accelerate) (0.4.6)
|
||||
Requirement already satisfied: MarkupSafe>=2.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from jinja2->torch>=1.10.0->accelerate) (3.0.2)
|
||||
Requirement already satisfied: charset-normalizer<4,>=2 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (3.4.0)
|
||||
Requirement already satisfied: idna<4,>=2.5 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (3.10)
|
||||
Requirement already satisfied: urllib3<3,>=1.21.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (2.2.3)
|
||||
Requirement already satisfied: certifi>=2017.4.17 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (2024.8.30)
|
||||
Downloading accelerate-1.0.1-py3-none-any.whl (330 kB)
|
||||
Downloading psutil-6.1.0-cp37-abi3-win_amd64.whl (254 kB)
|
||||
Installing collected packages: psutil, accelerate
|
||||
Successfully installed accelerate-1.0.1 psutil-6.1.0
|
||||
@@ -1,336 +0,0 @@
|
||||
import spaces
|
||||
import gradio as gr
|
||||
from huggingface_hub import InferenceClient
|
||||
from torch import nn
|
||||
from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
|
||||
from pathlib import Path
|
||||
import torch
|
||||
import torch.amp.autocast_mode
|
||||
from PIL import Image
|
||||
import os
|
||||
import torchvision.transforms.functional as TVF
|
||||
|
||||
|
||||
CLIP_PATH = "google/siglip-so400m-patch14-384"
|
||||
CHECKPOINT_PATH = Path("cgrkzexw-599808")
|
||||
TITLE = "<h1><center>JoyCaption Alpha Two (2024-09-26a)</center></h1>"
|
||||
CAPTION_TYPE_MAP = {
|
||||
"Descriptive": [
|
||||
"Write a descriptive caption for this image in a formal tone.",
|
||||
"Write a descriptive caption for this image in a formal tone within {word_count} words.",
|
||||
"Write a {length} descriptive caption for this image in a formal tone.",
|
||||
],
|
||||
"Descriptive (Informal)": [
|
||||
"Write a descriptive caption for this image in a casual tone.",
|
||||
"Write a descriptive caption for this image in a casual tone within {word_count} words.",
|
||||
"Write a {length} descriptive caption for this image in a casual tone.",
|
||||
],
|
||||
"Training Prompt": [
|
||||
"Write a stable diffusion prompt for this image.",
|
||||
"Write a stable diffusion prompt for this image within {word_count} words.",
|
||||
"Write a {length} stable diffusion prompt for this image.",
|
||||
],
|
||||
"MidJourney": [
|
||||
"Write a MidJourney prompt for this image.",
|
||||
"Write a MidJourney prompt for this image within {word_count} words.",
|
||||
"Write a {length} MidJourney prompt for this image.",
|
||||
],
|
||||
"Booru tag list": [
|
||||
"Write a list of Booru tags for this image.",
|
||||
"Write a list of Booru tags for this image within {word_count} words.",
|
||||
"Write a {length} list of Booru tags for this image.",
|
||||
],
|
||||
"Booru-like tag list": [
|
||||
"Write a list of Booru-like tags for this image.",
|
||||
"Write a list of Booru-like tags for this image within {word_count} words.",
|
||||
"Write a {length} list of Booru-like tags for this image.",
|
||||
],
|
||||
"Art Critic": [
|
||||
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc.",
|
||||
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it within {word_count} words.",
|
||||
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it {length}.",
|
||||
],
|
||||
"Product Listing": [
|
||||
"Write a caption for this image as though it were a product listing.",
|
||||
"Write a caption for this image as though it were a product listing. Keep it under {word_count} words.",
|
||||
"Write a {length} caption for this image as though it were a product listing.",
|
||||
],
|
||||
"Social Media Post": [
|
||||
"Write a caption for this image as if it were being used for a social media post.",
|
||||
"Write a caption for this image as if it were being used for a social media post. Limit the caption to {word_count} words.",
|
||||
"Write a {length} caption for this image as if it were being used for a social media post.",
|
||||
],
|
||||
}
|
||||
|
||||
HF_TOKEN = os.environ.get("HF_TOKEN", None)
|
||||
|
||||
|
||||
class ImageAdapter(nn.Module):
|
||||
def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
|
||||
super().__init__()
|
||||
self.deep_extract = deep_extract
|
||||
|
||||
if self.deep_extract:
|
||||
input_features = input_features * 5
|
||||
|
||||
self.linear1 = nn.Linear(input_features, output_features)
|
||||
self.activation = nn.GELU()
|
||||
self.linear2 = nn.Linear(output_features, output_features)
|
||||
self.ln1 = nn.Identity() if not ln1 else nn.LayerNorm(input_features)
|
||||
self.pos_emb = None if not pos_emb else nn.Parameter(torch.zeros(num_image_tokens, input_features))
|
||||
|
||||
# Other tokens (<|image_start|>, <|image_end|>, <|eot_id|>)
|
||||
self.other_tokens = nn.Embedding(3, output_features)
|
||||
self.other_tokens.weight.data.normal_(mean=0.0, std=0.02) # Matches HF's implementation of llama3
|
||||
|
||||
def forward(self, vision_outputs: torch.Tensor):
|
||||
if self.deep_extract:
|
||||
x = torch.concat((
|
||||
vision_outputs[-2],
|
||||
vision_outputs[3],
|
||||
vision_outputs[7],
|
||||
vision_outputs[13],
|
||||
vision_outputs[20],
|
||||
), dim=-1)
|
||||
assert len(x.shape) == 3, f"Expected 3, got {len(x.shape)}" # batch, tokens, features
|
||||
assert x.shape[-1] == vision_outputs[-2].shape[-1] * 5, f"Expected {vision_outputs[-2].shape[-1] * 5}, got {x.shape[-1]}"
|
||||
else:
|
||||
x = vision_outputs[-2]
|
||||
|
||||
x = self.ln1(x)
|
||||
|
||||
if self.pos_emb is not None:
|
||||
assert x.shape[-2:] == self.pos_emb.shape, f"Expected {self.pos_emb.shape}, got {x.shape[-2:]}"
|
||||
x = x + self.pos_emb
|
||||
|
||||
x = self.linear1(x)
|
||||
x = self.activation(x)
|
||||
x = self.linear2(x)
|
||||
|
||||
# <|image_start|>, IMAGE, <|image_end|>
|
||||
other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1))
|
||||
assert other_tokens.shape == (x.shape[0], 2, x.shape[2]), f"Expected {(x.shape[0], 2, x.shape[2])}, got {other_tokens.shape}"
|
||||
x = torch.cat((other_tokens[:, 0:1], x, other_tokens[:, 1:2]), dim=1)
|
||||
|
||||
return x
|
||||
|
||||
def get_eot_embedding(self):
|
||||
return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0)
|
||||
|
||||
|
||||
|
||||
# Load CLIP
|
||||
print("Loading CLIP")
|
||||
clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
|
||||
clip_model = AutoModel.from_pretrained(CLIP_PATH)
|
||||
clip_model = clip_model.vision_model
|
||||
|
||||
assert (CHECKPOINT_PATH / "clip_model.pt").exists()
|
||||
print("Loading VLM's custom vision model")
|
||||
checkpoint = torch.load(CHECKPOINT_PATH / "clip_model.pt", map_location='cpu')
|
||||
checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()}
|
||||
clip_model.load_state_dict(checkpoint)
|
||||
del checkpoint
|
||||
|
||||
clip_model.eval()
|
||||
clip_model.requires_grad_(False)
|
||||
clip_model.to("cuda")
|
||||
|
||||
|
||||
# Tokenizer
|
||||
print("Loading tokenizer")
|
||||
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_PATH / "text_model", use_fast=True)
|
||||
assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
|
||||
|
||||
# LLM
|
||||
print("Loading LLM")
|
||||
print("Loading VLM's custom text model")
|
||||
text_model = AutoModelForCausalLM.from_pretrained(CHECKPOINT_PATH / "text_model", device_map=0, torch_dtype=torch.bfloat16)
|
||||
text_model.eval()
|
||||
|
||||
# Image Adapter
|
||||
print("Loading image adapter")
|
||||
image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False)
|
||||
image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu"))
|
||||
image_adapter.eval()
|
||||
image_adapter.to("cuda")
|
||||
|
||||
|
||||
@spaces.GPU()
|
||||
@torch.no_grad()
|
||||
def stream_chat(input_image: Image.Image, caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str, custom_prompt: str) -> tuple[str, str]:
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# 'any' means no length specified
|
||||
length = None if caption_length == "any" else caption_length
|
||||
|
||||
if isinstance(length, str):
|
||||
try:
|
||||
length = int(length)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Build prompt
|
||||
if length is None:
|
||||
map_idx = 0
|
||||
elif isinstance(length, int):
|
||||
map_idx = 1
|
||||
elif isinstance(length, str):
|
||||
map_idx = 2
|
||||
else:
|
||||
raise ValueError(f"Invalid caption length: {length}")
|
||||
|
||||
prompt_str = CAPTION_TYPE_MAP[caption_type][map_idx]
|
||||
|
||||
# Add extra options
|
||||
if len(extra_options) > 0:
|
||||
prompt_str += " " + " ".join(extra_options)
|
||||
|
||||
# Add name, length, word_count
|
||||
prompt_str = prompt_str.format(name=name_input, length=caption_length, word_count=caption_length)
|
||||
|
||||
if custom_prompt.strip() != "":
|
||||
prompt_str = custom_prompt.strip()
|
||||
|
||||
# For debugging
|
||||
print(f"Prompt: {prompt_str}")
|
||||
|
||||
# Preprocess image
|
||||
# NOTE: I found the default processor for so400M to have worse results than just using PIL directly
|
||||
#image = clip_processor(images=input_image, return_tensors='pt').pixel_values
|
||||
image = input_image.resize((384, 384), Image.LANCZOS)
|
||||
pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
|
||||
pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
|
||||
pixel_values = pixel_values.to('cuda')
|
||||
|
||||
# Embed image
|
||||
# This results in Batch x Image Tokens x Features
|
||||
with torch.amp.autocast_mode.autocast('cuda', enabled=True):
|
||||
vision_outputs = clip_model(pixel_values=pixel_values, output_hidden_states=True)
|
||||
embedded_images = image_adapter(vision_outputs.hidden_states)
|
||||
embedded_images = embedded_images.to('cuda')
|
||||
|
||||
# Build the conversation
|
||||
convo = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful image captioner.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt_str,
|
||||
},
|
||||
]
|
||||
|
||||
# Format the conversation
|
||||
convo_string = tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = True)
|
||||
assert isinstance(convo_string, str)
|
||||
|
||||
# Tokenize the conversation
|
||||
# prompt_str is tokenized separately so we can do the calculations below
|
||||
convo_tokens = tokenizer.encode(convo_string, return_tensors="pt", add_special_tokens=False, truncation=False)
|
||||
prompt_tokens = tokenizer.encode(prompt_str, return_tensors="pt", add_special_tokens=False, truncation=False)
|
||||
assert isinstance(convo_tokens, torch.Tensor) and isinstance(prompt_tokens, torch.Tensor)
|
||||
convo_tokens = convo_tokens.squeeze(0) # Squeeze just to make the following easier
|
||||
prompt_tokens = prompt_tokens.squeeze(0)
|
||||
|
||||
# Calculate where to inject the image
|
||||
eot_id_indices = (convo_tokens == tokenizer.convert_tokens_to_ids("<|eot_id|>")).nonzero(as_tuple=True)[0].tolist()
|
||||
assert len(eot_id_indices) == 2, f"Expected 2 <|eot_id|> tokens, got {len(eot_id_indices)}"
|
||||
|
||||
preamble_len = eot_id_indices[1] - prompt_tokens.shape[0] # Number of tokens before the prompt
|
||||
|
||||
# Embed the tokens
|
||||
convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to('cuda'))
|
||||
|
||||
# Construct the input
|
||||
input_embeds = torch.cat([
|
||||
convo_embeds[:, :preamble_len], # Part before the prompt
|
||||
embedded_images.to(dtype=convo_embeds.dtype), # Image
|
||||
convo_embeds[:, preamble_len:], # The prompt and anything after it
|
||||
], dim=1).to('cuda')
|
||||
|
||||
input_ids = torch.cat([
|
||||
convo_tokens[:preamble_len].unsqueeze(0),
|
||||
torch.zeros((1, embedded_images.shape[1]), dtype=torch.long), # Dummy tokens for the image (TODO: Should probably use a special token here so as not to confuse any generation algorithms that might be inspecting the input)
|
||||
convo_tokens[preamble_len:].unsqueeze(0),
|
||||
], dim=1).to('cuda')
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
# Debugging
|
||||
print(f"Input to model: {repr(tokenizer.decode(input_ids[0]))}")
|
||||
|
||||
#generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None)
|
||||
#generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, suppress_tokens=None)
|
||||
generate_ids = text_model.generate(input_ids, inputs_embeds=input_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, suppress_tokens=None) # Uses the default which is temp=0.6, top_p=0.9
|
||||
|
||||
# Trim off the prompt
|
||||
generate_ids = generate_ids[:, input_ids.shape[1]:]
|
||||
if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids("<|eot_id|>"):
|
||||
generate_ids = generate_ids[:, :-1]
|
||||
|
||||
caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]
|
||||
|
||||
return prompt_str, caption.strip()
|
||||
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
gr.HTML(TITLE)
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
input_image = gr.Image(type="pil", label="Input Image")
|
||||
|
||||
caption_type = gr.Dropdown(
|
||||
choices=["Descriptive", "Descriptive (Informal)", "Training Prompt", "MidJourney", "Booru tag list", "Booru-like tag list", "Art Critic", "Product Listing", "Social Media Post"],
|
||||
label="Caption Type",
|
||||
value="Descriptive",
|
||||
)
|
||||
|
||||
caption_length = gr.Dropdown(
|
||||
choices=["any", "very short", "short", "medium-length", "long", "very long"] +
|
||||
[str(i) for i in range(20, 261, 10)],
|
||||
label="Caption Length",
|
||||
value="long",
|
||||
)
|
||||
|
||||
extra_options = gr.CheckboxGroup(
|
||||
choices=[
|
||||
"If there is a person/character in the image you must refer to them as {name}.",
|
||||
"Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).",
|
||||
"Include information about lighting.",
|
||||
"Include information about camera angle.",
|
||||
"Include information about whether there is a watermark or not.",
|
||||
"Include information about whether there are JPEG artifacts or not.",
|
||||
"If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.",
|
||||
"Do NOT include anything sexual; keep it PG.",
|
||||
"Do NOT mention the image's resolution.",
|
||||
"You MUST include information about the subjective aesthetic quality of the image from low to very high.",
|
||||
"Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.",
|
||||
"Do NOT mention any text that is in the image.",
|
||||
"Specify the depth of field and whether the background is in focus or blurred.",
|
||||
"If applicable, mention the likely use of artificial or natural lighting sources.",
|
||||
"Do NOT use any ambiguous language.",
|
||||
"Include whether the image is sfw, suggestive, or nsfw.",
|
||||
"ONLY describe the most important elements of the image."
|
||||
],
|
||||
label="Extra Options"
|
||||
)
|
||||
|
||||
name_input = gr.Textbox(label="Person/Character Name (if applicable)")
|
||||
gr.Markdown("**Note:** Name input is only used if an Extra Option is selected that requires it.")
|
||||
|
||||
custom_prompt = gr.Textbox(label="Custom Prompt (optional, will override all other settings)")
|
||||
gr.Markdown("**Note:** Alpha Two is not a general instruction follower and will not follow prompts outside its training data well. Use this feature with caution.")
|
||||
|
||||
run_button = gr.Button("Caption")
|
||||
|
||||
with gr.Column():
|
||||
output_prompt = gr.Textbox(label="Prompt that was used")
|
||||
output_caption = gr.Textbox(label="Caption")
|
||||
|
||||
run_button.click(fn=stream_chat, inputs=[input_image, caption_type, caption_length, extra_options, name_input, custom_prompt], outputs=[output_prompt, output_caption])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo.launch()
|
||||
@@ -1,423 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Use JoyCaption to caption images.
|
||||
"""
|
||||
import argparse
|
||||
import dataclasses
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torch.amp
|
||||
import torchvision.transforms.functional as TVF
|
||||
from PIL import Image
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from tqdm import tqdm
|
||||
from transformers import (
|
||||
AutoTokenizer,
|
||||
LlavaForConditionalGeneration,
|
||||
PreTrainedTokenizer,
|
||||
PreTrainedTokenizerFast,
|
||||
)
|
||||
|
||||
|
||||
def none_or_type(value, desired_type):
|
||||
if value == "None":
|
||||
return None
|
||||
return desired_type(value)
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--glob", type=str, help="Glob pattern to find images")
|
||||
parser.add_argument("--filelist", type=str, help="File containing list of images")
|
||||
parser.add_argument("--prompt", type=str, help="Prompt to use")
|
||||
parser.add_argument(
|
||||
"--prompt-file", type=str, help="JSON file containing prompts to use"
|
||||
)
|
||||
parser.add_argument("--batch-size", type=int, default=1, help="Batch size")
|
||||
parser.add_argument(
|
||||
"--greedy", action="store_true", help="Use greedy decoding instead of sampling"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--temperature", type=float, default=0.6, help="Sampling temperature"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-p", type=lambda x: none_or_type(x, float), default=0.9, help="Top-p sampling"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-k", type=lambda x: none_or_type(x, int), default=None, help="Top-k sampling"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=256,
|
||||
help="Maximum length of the generated caption (in tokens)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-workers",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Number of workers loading images in parallel",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="fancyfeast/llama-joycaption-alpha-two-hf-llava",
|
||||
help="Model to use",
|
||||
)
|
||||
|
||||
|
||||
PIL.Image.MAX_IMAGE_PIXELS = 933120000 # Quiets Pillow from giving warnings on really large images (WARNING: Exposes a risk of DoS from malicious images)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Prompt:
|
||||
prompt: str
|
||||
weight: float
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def main():
|
||||
# Logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s"
|
||||
)
|
||||
|
||||
# Parse arguments
|
||||
args = parser.parse_args()
|
||||
logging.info(f"Arguments: {args}")
|
||||
|
||||
args.prompt = "Please describe the image."
|
||||
# Make sure we have a prompt or a prompt file
|
||||
prompts = parse_prompts(args.prompt, args.prompt_file)
|
||||
|
||||
args.filelist = (
|
||||
"C:\\Users\\27698\\Desktop\\node\\12\\12.txt"
|
||||
)
|
||||
# Find the images
|
||||
image_paths = find_images(args.glob, args.filelist)
|
||||
if len(image_paths) == 0:
|
||||
logging.warning("No images found")
|
||||
return
|
||||
logging.info(f"Found {len(image_paths)} images")
|
||||
|
||||
# Ignore all images that already have captions
|
||||
image_paths = [
|
||||
path for path in image_paths if not Path(path).with_suffix(".txt").exists()
|
||||
]
|
||||
|
||||
# Load JoyCaption
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=True)
|
||||
assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(
|
||||
tokenizer, PreTrainedTokenizerFast
|
||||
), f"Tokenizer is of type {type(tokenizer)}"
|
||||
llava_model = LlavaForConditionalGeneration.from_pretrained(
|
||||
args.model, torch_dtype="bfloat16"
|
||||
)
|
||||
assert isinstance(llava_model, LlavaForConditionalGeneration)
|
||||
|
||||
dataset = ImageDataset(
|
||||
prompts,
|
||||
image_paths,
|
||||
tokenizer,
|
||||
llava_model.config.image_token_index,
|
||||
llava_model.config.image_seq_length,
|
||||
)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
collate_fn=dataset.collate_fn,
|
||||
num_workers=args.num_workers,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
batch_size=args.batch_size,
|
||||
)
|
||||
end_of_header_id = tokenizer.convert_tokens_to_ids("<|end_header_id|>")
|
||||
end_of_turn_id = tokenizer.convert_tokens_to_ids("<|eot_id|>")
|
||||
assert isinstance(end_of_header_id, int) and isinstance(end_of_turn_id, int)
|
||||
|
||||
pbar = tqdm(total=len(image_paths), desc="Captioning images...", dynamic_ncols=True)
|
||||
for batch in dataloader:
|
||||
vision_dtype = (
|
||||
llava_model.vision_tower.vision_model.embeddings.patch_embedding.weight.dtype
|
||||
)
|
||||
vision_device = (
|
||||
llava_model.vision_tower.vision_model.embeddings.patch_embedding.weight.device
|
||||
)
|
||||
language_device = (
|
||||
llava_model.language_model.get_input_embeddings().weight.device
|
||||
)
|
||||
|
||||
# Move to GPU
|
||||
pixel_values = batch["pixel_values"].to(vision_device, non_blocking=True)
|
||||
input_ids = batch["input_ids"].to(language_device, non_blocking=True)
|
||||
attention_mask = batch["attention_mask"].to(language_device, non_blocking=True)
|
||||
|
||||
# Normalize the image
|
||||
pixel_values = pixel_values / 255.0
|
||||
pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
|
||||
pixel_values = pixel_values.to(vision_dtype)
|
||||
|
||||
# Generate the captions
|
||||
generate_ids = llava_model.generate(
|
||||
input_ids=input_ids,
|
||||
pixel_values=pixel_values,
|
||||
attention_mask=attention_mask,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=not args.greedy,
|
||||
suppress_tokens=None,
|
||||
use_cache=True,
|
||||
temperature=args.temperature,
|
||||
top_k=args.top_k,
|
||||
top_p=args.top_p,
|
||||
)
|
||||
|
||||
# Trim off the prompts
|
||||
assert isinstance(generate_ids, torch.Tensor)
|
||||
generate_ids = generate_ids.tolist()
|
||||
generate_ids = [
|
||||
trim_off_prompt(ids, end_of_header_id, end_of_turn_id)
|
||||
for ids in generate_ids
|
||||
]
|
||||
|
||||
# Decode the captions
|
||||
captions = tokenizer.batch_decode(
|
||||
generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False
|
||||
)
|
||||
captions = [c.strip() for c in captions]
|
||||
|
||||
for path, caption in zip(batch["paths"], captions):
|
||||
write_caption(Path(path), caption)
|
||||
|
||||
pbar.update(len(captions))
|
||||
|
||||
|
||||
def trim_off_prompt(input_ids: list[int], eoh_id: int, eot_id: int) -> list[int]:
|
||||
# Trim off the prompt
|
||||
while True:
|
||||
try:
|
||||
i = input_ids.index(eoh_id)
|
||||
except ValueError:
|
||||
break
|
||||
|
||||
input_ids = input_ids[i + 1 :]
|
||||
|
||||
# Trim off the end
|
||||
try:
|
||||
i = input_ids.index(eot_id)
|
||||
except ValueError:
|
||||
return input_ids
|
||||
|
||||
return input_ids[:i]
|
||||
|
||||
|
||||
def write_caption(image_path: Path, caption: str):
|
||||
caption_path = image_path.with_suffix(".txt")
|
||||
|
||||
try:
|
||||
f = os.open(
|
||||
caption_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL
|
||||
) # Write-only, create if not exist, fail if exists
|
||||
except FileExistsError:
|
||||
logging.warning(f"Caption file '{caption_path}' already exists")
|
||||
return
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to open caption file '{caption_path}': {e}")
|
||||
return
|
||||
|
||||
try:
|
||||
os.write(f, caption.encode("utf-8"))
|
||||
os.close(f)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to write caption to '{caption_path}': {e}")
|
||||
return
|
||||
|
||||
|
||||
class ImageDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
prompts: list[Prompt],
|
||||
paths: list[Path],
|
||||
tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
|
||||
image_token_id: int,
|
||||
image_seq_length: int,
|
||||
):
|
||||
self.prompts = prompts
|
||||
self.paths = paths
|
||||
self.tokenizer = tokenizer
|
||||
self.image_token_id = image_token_id
|
||||
self.image_seq_length = image_seq_length
|
||||
self.pad_token_id = tokenizer.pad_token_id
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
path = self.paths[idx]
|
||||
|
||||
# Pick a prompt
|
||||
prompt_str = random.choices(
|
||||
self.prompts, weights=[p.weight for p in self.prompts]
|
||||
)[0].prompt
|
||||
|
||||
# Preprocess image
|
||||
# NOTE: I don't use the Processor here and instead do it manually.
|
||||
# This is because in my testing a simple resize in Pillow yields higher quality results than the Processor,
|
||||
# and the Processor had some buggy behavior on some images.
|
||||
# And yes, with the so400m model, the model expects the image to be squished into a square, not padded.
|
||||
try:
|
||||
image = Image.open(path)
|
||||
if image.size != (384, 384):
|
||||
image = image.resize((384, 384), Image.LANCZOS)
|
||||
image = image.convert("RGB")
|
||||
pixel_values = TVF.pil_to_tensor(image)
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to load image '{path}': {e}")
|
||||
pixel_values = None # Will be filtered out later
|
||||
|
||||
# Build the conversation
|
||||
convo = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful image captioner.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt_str,
|
||||
},
|
||||
]
|
||||
|
||||
# Format the conversation
|
||||
convo_string = self.tokenizer.apply_chat_template(
|
||||
convo, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
assert isinstance(convo_string, str)
|
||||
|
||||
# Tokenize the conversation
|
||||
convo_tokens = self.tokenizer.encode(
|
||||
convo_string, add_special_tokens=False, truncation=False
|
||||
)
|
||||
|
||||
# Repeat the image tokens
|
||||
input_tokens = []
|
||||
for token in convo_tokens:
|
||||
if token == self.image_token_id:
|
||||
input_tokens.extend([self.image_token_id] * self.image_seq_length)
|
||||
else:
|
||||
input_tokens.append(token)
|
||||
|
||||
input_ids = torch.tensor(input_tokens, dtype=torch.long)
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
return {
|
||||
"path": path,
|
||||
"pixel_values": pixel_values,
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
}
|
||||
|
||||
def collate_fn(self, batch: list[dict]) -> dict:
|
||||
# Filter out images that failed to load
|
||||
batch = [item for item in batch if item["pixel_values"] is not None]
|
||||
|
||||
# Pad input_ids and attention_mask
|
||||
# Have to use left padding because HF's generate can't handle right padding it seems
|
||||
max_length = max(item["input_ids"].shape[0] for item in batch)
|
||||
n_pad = [max_length - item["input_ids"].shape[0] for item in batch]
|
||||
input_ids = torch.stack(
|
||||
[
|
||||
torch.nn.functional.pad(
|
||||
item["input_ids"], (n, 0), value=self.pad_token_id
|
||||
)
|
||||
for item, n in zip(batch, n_pad)
|
||||
]
|
||||
)
|
||||
attention_mask = torch.stack(
|
||||
[
|
||||
torch.nn.functional.pad(item["attention_mask"], (n, 0), value=0)
|
||||
for item, n in zip(batch, n_pad)
|
||||
]
|
||||
)
|
||||
|
||||
# Stack pixel values
|
||||
pixel_values = torch.stack([item["pixel_values"] for item in batch])
|
||||
|
||||
# Paths
|
||||
paths = [item["path"] for item in batch]
|
||||
|
||||
return {
|
||||
"paths": paths,
|
||||
"pixel_values": pixel_values,
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
}
|
||||
|
||||
|
||||
def parse_prompts(prompt_str: str | None, prompt_file: str | None) -> list[Prompt]:
|
||||
if prompt_str is not None and prompt_file is not None:
|
||||
raise ValueError("Cannot specify both --prompt and --prompt-file")
|
||||
|
||||
if prompt_str is not None:
|
||||
return [Prompt(prompt=prompt_str, weight=1.0)]
|
||||
|
||||
if prompt_file is None:
|
||||
raise ValueError("Must specify either --prompt or --prompt-file")
|
||||
|
||||
data = json.loads(Path(prompt_file).read_text())
|
||||
|
||||
if not isinstance(data, list):
|
||||
raise ValueError("Expected JSON file to contain a list of prompts")
|
||||
|
||||
prompts = []
|
||||
|
||||
for item in data:
|
||||
if isinstance(item, str):
|
||||
prompts.append(Prompt(prompt=item, weight=1.0))
|
||||
elif (
|
||||
isinstance(item, dict)
|
||||
and "prompt" in item
|
||||
and "weight" in item
|
||||
and isinstance(item["prompt"], str)
|
||||
and isinstance(item["weight"], (int, float))
|
||||
):
|
||||
prompts.append(Prompt(prompt=item["prompt"], weight=item["weight"]))
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid prompt in JSON file. Should be either a string or an object with 'prompt' and 'weight' fields: {item}"
|
||||
)
|
||||
|
||||
if len(prompts) == 0:
|
||||
raise ValueError("No prompts found in JSON file")
|
||||
|
||||
if sum(p.weight for p in prompts) <= 0.0:
|
||||
raise ValueError("Prompt weights must sum to a positive number")
|
||||
|
||||
return prompts
|
||||
|
||||
|
||||
def find_images(glob: str | None, filelist: str | Path | None) -> list[Path]:
|
||||
if glob is None and filelist is None:
|
||||
raise ValueError("Must specify either --glob or --filelist")
|
||||
|
||||
paths = []
|
||||
|
||||
if glob is not None:
|
||||
paths.extend(Path(".").glob(glob))
|
||||
|
||||
if filelist is not None:
|
||||
paths.extend(
|
||||
(
|
||||
Path(line.strip())
|
||||
for line in Path(filelist).read_text().strip().splitlines()
|
||||
if line.strip() != ""
|
||||
)
|
||||
)
|
||||
|
||||
return paths
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,75 +0,0 @@
|
||||
import torch
|
||||
import torch.amp
|
||||
import torchvision.transforms.functional as TVF
|
||||
from PIL import Image
|
||||
from transformers import AutoTokenizer, LlavaForConditionalGeneration
|
||||
|
||||
|
||||
IMAGE_PATH = "C:/Users/27698/Desktop/node/12/00001.png"
|
||||
PROMPT = "Write a long descriptive caption for this image in a formal tone."
|
||||
MODEL_NAME = "fancyfeast/llama-joycaption-alpha-two-hf-llava"
|
||||
|
||||
|
||||
# Load JoyCaption
|
||||
# bfloat16 is the native dtype of the LLM used in JoyCaption (Llama 3.1)
|
||||
# device_map=0 loads the model into the first GPU
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
|
||||
llava_model = LlavaForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype="bfloat16", device_map="cuda:0")
|
||||
llava_model.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
# Load and preprocess image
|
||||
# Normally you would use the Processor here, but the image module's processor
|
||||
# has some buggy behavior and a simple resize in Pillow yields higher quality results
|
||||
image = Image.open(IMAGE_PATH)
|
||||
|
||||
if image.size != (384, 384):
|
||||
image = image.resize((384, 384), Image.LANCZOS)
|
||||
|
||||
image = image.convert("RGB")
|
||||
pixel_values = TVF.pil_to_tensor(image)
|
||||
|
||||
# Normalize the image
|
||||
pixel_values = pixel_values / 255.0
|
||||
pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
|
||||
pixel_values = pixel_values.to(torch.bfloat16).unsqueeze(0)
|
||||
|
||||
# Build the conversation
|
||||
convo = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful image captioner.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": PROMPT,
|
||||
},
|
||||
]
|
||||
|
||||
# Format the conversation
|
||||
convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
|
||||
|
||||
# Tokenize the conversation
|
||||
convo_tokens = tokenizer.encode(convo_string, add_special_tokens=False, truncation=False)
|
||||
|
||||
# Repeat the image tokens
|
||||
input_tokens = []
|
||||
for token in convo_tokens:
|
||||
if token == llava_model.config.image_token_index:
|
||||
input_tokens.extend([llava_model.config.image_token_index] * llava_model.config.image_seq_length)
|
||||
else:
|
||||
input_tokens.append(token)
|
||||
|
||||
input_ids = torch.tensor(input_tokens, dtype=torch.long).unsqueeze(0)
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
# Generate the caption
|
||||
generate_ids = llava_model.generate(input_ids=input_ids.to('cuda'), pixel_values=pixel_values.to('cuda'), attention_mask=attention_mask.to('cuda'), max_new_tokens=300, do_sample=True, suppress_tokens=None, use_cache=True)[0]
|
||||
|
||||
# Trim off the prompt
|
||||
generate_ids = generate_ids[input_ids.shape[1]:]
|
||||
|
||||
# Decode the caption
|
||||
caption = tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
||||
caption = caption.strip()
|
||||
print(caption)
|
||||
@@ -1,2 +0,0 @@
|
||||
@echo off
|
||||
pyinstaller -F --upx-dir="C:\\Users\\27698\\Desktop\\upx-4.2.4-win64\upx.exe" lama_inpaint.py
|
||||
@@ -1,173 +0,0 @@
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
from typing import Union
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
|
||||
|
||||
# 判断sys.argv 的长度，如果小于2，说明没有传入参数，设置初始参数
|
||||
# if len(sys.argv) < 2:
|
||||
# sys.argv = [
|
||||
# "C:/Users/27698/Desktop/LAITool/resources/scripts/lama/lama_inpaint.exe",
|
||||
# "-l",
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\mjTest\\data\\mask\\temp\\1717508661218.png",
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\mjTest\\data\\mask\\mask_temp_1717508662659.png",
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\mjTest\\data\\mask\\temp\\1717508564042.png",
|
||||
# ]
|
||||
print(sys.argv)
|
||||
|
||||
if getattr(sys, "frozen", False):
|
||||
cript_directory = os.path.dirname(sys.executable)
|
||||
elif __file__:
|
||||
cript_directory = os.path.dirname(__file__)
|
||||
|
||||
link_name = os.path.join(os.path.expanduser("~"), "big_lama.pt")
|
||||
cu_name = os.path.join(cript_directory, "model\\big-lama.pt")
|
||||
mode_pa = link_name
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
# # 判断model_path是否存在，如果不存在，设置默认值
|
||||
if not os.path.exists(link_name):
|
||||
os.system(f'mklink "{link_name}" "{cu_name}"')
|
||||
print("Params: <runtime-config.json>")
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def get_image(image):
|
||||
if isinstance(image, Image.Image):
|
||||
img = np.array(image)
|
||||
elif isinstance(image, np.ndarray):
|
||||
img = image.copy()
|
||||
else:
|
||||
raise Exception("Input image should be either PIL Image or numpy array!")
|
||||
|
||||
if img.ndim == 3:
|
||||
img = np.transpose(img, (2, 0, 1)) # chw
|
||||
elif img.ndim == 2:
|
||||
img = img[np.newaxis, ...]
|
||||
|
||||
assert img.ndim == 3
|
||||
|
||||
img = img.astype(np.float32) / 255
|
||||
return img
|
||||
|
||||
|
||||
def ceil_modulo(x, mod):
|
||||
if x % mod == 0:
|
||||
return x
|
||||
return (x // mod + 1) * mod
|
||||
|
||||
|
||||
def scale_image(img, factor, interpolation=cv2.INTER_AREA):
|
||||
if img.shape[0] == 1:
|
||||
img = img[0]
|
||||
else:
|
||||
img = np.transpose(img, (1, 2, 0))
|
||||
|
||||
img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation)
|
||||
|
||||
if img.ndim == 2:
|
||||
img = img[None, ...]
|
||||
else:
|
||||
img = np.transpose(img, (2, 0, 1))
|
||||
return img
|
||||
|
||||
|
||||
def pad_img_to_modulo(img, mod):
|
||||
channels, height, width = img.shape
|
||||
out_height = ceil_modulo(height, mod)
|
||||
out_width = ceil_modulo(width, mod)
|
||||
return np.pad(
|
||||
img,
|
||||
((0, 0), (0, out_height - height), (0, out_width - width)),
|
||||
mode="symmetric",
|
||||
)
|
||||
|
||||
|
||||
def prepare_img_and_mask(image, mask, device, pad_out_to_modulo=8, scale_factor=None):
|
||||
out_image = get_image(image)
|
||||
out_mask = get_image(mask)
|
||||
|
||||
if scale_factor is not None:
|
||||
out_image = scale_image(out_image, 1)
|
||||
out_mask = scale_image(out_mask, scale_factor, interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
if pad_out_to_modulo is not None and pad_out_to_modulo > 1:
|
||||
out_image = pad_img_to_modulo(out_image, pad_out_to_modulo)
|
||||
out_mask = pad_img_to_modulo(out_mask, pad_out_to_modulo)
|
||||
|
||||
out_image = torch.from_numpy(out_image).unsqueeze(0).to(device)
|
||||
out_mask = torch.from_numpy(out_mask).unsqueeze(0).to(device)
|
||||
|
||||
out_mask = (out_mask > 0) * 1
|
||||
|
||||
return out_image, out_mask
|
||||
|
||||
|
||||
class LamaInpaint:
|
||||
def __init__(
|
||||
self,
|
||||
device,
|
||||
model_path=None,
|
||||
) -> None:
|
||||
if model_path is None:
|
||||
model_path = os.path.join(cript_directory, "model\\big-lama.pt")
|
||||
|
||||
self.model = torch.jit.load(model_path, map_location=device)
|
||||
self.model.eval()
|
||||
self.model.to(device)
|
||||
self.device = device
|
||||
|
||||
def run(
|
||||
self,
|
||||
image: Union[Image.Image, np.ndarray],
|
||||
mask: Union[Image.Image, np.ndarray],
|
||||
):
|
||||
if isinstance(image, np.ndarray):
|
||||
orig_height, orig_width = image.shape[:2]
|
||||
else:
|
||||
orig_height, orig_width = np.array(image).shape[:2]
|
||||
|
||||
# image_width = image.shape[1]
|
||||
# mask_width = mask.shape[1]
|
||||
scale = image.width / mask.width
|
||||
image, mask = prepare_img_and_mask(image, mask, self.device, 8, scale)
|
||||
with torch.inference_mode():
|
||||
inpainted = self.model(image, mask)
|
||||
cur_res = inpainted[0].permute(1, 2, 0).detach().cpu().numpy()
|
||||
cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
|
||||
cur_res = cur_res[:orig_height, :orig_width]
|
||||
return cur_res
|
||||
|
||||
|
||||
try:
|
||||
de = "cpu"
|
||||
if torch.cuda.is_available():
|
||||
de = "cuda"
|
||||
|
||||
lama = LamaInpaint(de, mode_pa)
|
||||
|
||||
image_path = sys.argv[2]
|
||||
mask_path = sys.argv[3]
|
||||
output_path = sys.argv[4]
|
||||
|
||||
# 若是没有传递mask_path，需要自己计算mask区域
|
||||
# 使用Image.open打开图片
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
mask = Image.open(mask_path).convert("L")
|
||||
|
||||
res = lama.run(image, mask)
|
||||
# 将修复后的图片保存到本地
|
||||
img = Image.fromarray(res)
|
||||
# 使用 save 方法将图像保存到文件
|
||||
img.save(output_path)
|
||||
sys.exit(0)
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
sys.exit(str(e))
|
||||
@@ -1,43 +0,0 @@
|
||||
# -*- mode: python ; coding: utf-8 -*-
|
||||
|
||||
|
||||
a = Analysis(
|
||||
['lama_inpaint.py'],
|
||||
pathex=[],
|
||||
binaries=[],
|
||||
datas=[],
|
||||
hiddenimports=[],
|
||||
hookspath=[],
|
||||
hooksconfig={},
|
||||
runtime_hooks=[],
|
||||
excludes=[],
|
||||
noarchive=False,
|
||||
)
|
||||
pyz = PYZ(a.pure)
|
||||
|
||||
exe = EXE(
|
||||
pyz,
|
||||
a.scripts,
|
||||
[],
|
||||
exclude_binaries=True,
|
||||
name='lama_inpaint',
|
||||
debug=False,
|
||||
bootloader_ignore_signals=False,
|
||||
strip=False,
|
||||
upx=True,
|
||||
console=True,
|
||||
disable_windowed_traceback=False,
|
||||
argv_emulation=False,
|
||||
target_arch=None,
|
||||
codesign_identity=None,
|
||||
entitlements_file=None,
|
||||
)
|
||||
coll = COLLECT(
|
||||
exe,
|
||||
a.binaries,
|
||||
a.datas,
|
||||
strip=False,
|
||||
upx=True,
|
||||
upx_exclude=[],
|
||||
name='lama_inpaint',
|
||||
)
|
||||
@@ -1,2 +0,0 @@
|
||||
@echo off
|
||||
pyinstaller --upx-dir="C:\\Users\\27698\\Desktop\\upx-4.2.4-win64\upx.exe" local_whisper.py
|
||||
@@ -1,170 +0,0 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import public_tools
|
||||
from pathlib import Path
|
||||
from huggingface_hub import hf_hub_download
|
||||
from faster_whisper import WhisperModel
|
||||
|
||||
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
|
||||
|
||||
# 判断sys.argv 的长度，如果小于2，说明没有传入参数，设置初始参数
|
||||
# "C:\\Users\\27698\\Desktop\\LAITool\\resources\\scripts\\Lai.exe" -c "D:/来推项目集/7.4/娱乐：江湖大哥退休，去拍电影/scripts/output_crop_00001.json" "NVIDIA"
|
||||
# if len(sys.argv) < 2:
|
||||
# sys.argv = [
|
||||
# "C:\\Users\\27698\\Desktop\\LAITool\\resources\\scripts\\Lai.exe",
|
||||
# "-w",
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\test\\mjTestoutput_crop_00001.mp4",
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\test\data\\frame",
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\test\\tmp\\input_crop",
|
||||
# 30,
|
||||
# "NVIDIA",
|
||||
# ]
|
||||
|
||||
print(sys.argv)
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
print("Params: <runtime-config.json>")
|
||||
exit(0)
|
||||
|
||||
if getattr(sys, "frozen", False):
|
||||
cript_directory = os.path.dirname(sys.executable)
|
||||
elif __file__:
|
||||
cript_directory = os.path.dirname(__file__)
|
||||
|
||||
|
||||
def GetText(out_folder, mp3_folder):
|
||||
text = []
|
||||
# 先获取模型
|
||||
print("正在下载或加载模型")
|
||||
sys.stdout.flush()
|
||||
model_path = Path(
|
||||
hf_hub_download(repo_id="Systran/faster-whisper-large-v3", filename="model.bin")
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="preprocessor_config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="tokenizer.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="vocabulary.json",
|
||||
)
|
||||
model = WhisperModel(
|
||||
model_size_or_path=os.path.dirname(model_path),
|
||||
device="auto",
|
||||
local_files_only=True,
|
||||
)
|
||||
print("模型加载成功，开始识别")
|
||||
sys.stdout.flush()
|
||||
# 拿到指定文件夹里面的所有的MP3文件
|
||||
mp3_list = []
|
||||
for root, dirs, files in os.walk(mp3_folder):
|
||||
for file in files:
|
||||
if file.endswith(".mp3"):
|
||||
mp3_list.append(os.path.join(root, file))
|
||||
|
||||
for mp in mp3_list:
|
||||
segments, info = model.transcribe(
|
||||
mp,
|
||||
beam_size=5,
|
||||
language="zh",
|
||||
vad_filter=True,
|
||||
vad_parameters=dict(min_silence_duration_ms=1000),
|
||||
)
|
||||
tmp_text = ""
|
||||
for segment in segments:
|
||||
tmp_text += segment.text + "。"
|
||||
print(mp + "识别完成")
|
||||
sys.stdout.flush()
|
||||
text.append(tmp_text)
|
||||
|
||||
# 数据写出
|
||||
print("文本全部识别成功，正在写出")
|
||||
sys.stdout.flush()
|
||||
tools = public_tools.PublicTools()
|
||||
tools.write_to_file(text, os.path.join(out_folder, "文案.txt"))
|
||||
print("写出完成")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def GetTextTask(out_folder, mp, name):
|
||||
text = []
|
||||
# 先获取模型
|
||||
print("正在下载或加载模型")
|
||||
sys.stdout.flush()
|
||||
model_path = Path(
|
||||
hf_hub_download(repo_id="Systran/faster-whisper-large-v3", filename="model.bin")
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="preprocessor_config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="tokenizer.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="vocabulary.json",
|
||||
)
|
||||
model = WhisperModel(
|
||||
model_size_or_path=os.path.dirname(model_path),
|
||||
device="auto",
|
||||
local_files_only=True,
|
||||
)
|
||||
print("模型加载成功，开始识别")
|
||||
sys.stdout.flush()
|
||||
segments, info = model.transcribe(
|
||||
mp,
|
||||
beam_size=5,
|
||||
language="zh",
|
||||
vad_filter=True,
|
||||
vad_parameters=dict(min_silence_duration_ms=1000),
|
||||
)
|
||||
tmp_text = ""
|
||||
for segment in segments:
|
||||
tmp_text += segment.text + "。"
|
||||
print(mp + "识别完成")
|
||||
sys.stdout.flush()
|
||||
text.append(tmp_text)
|
||||
|
||||
# 数据写出
|
||||
sys.stdout.flush()
|
||||
tools = public_tools.PublicTools()
|
||||
tools.write_to_file(text, os.path.join(out_folder, name + ".txt"))
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
# GetTextTask(
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\mjTest",
|
||||
# "C:\\Users\\27698\\Desktop\\测试\\mjTest\\data\\frame\\00001.mp4",
|
||||
# "00001",
|
||||
# )
|
||||
|
||||
if sys.argv[1] == "-ts":
|
||||
GetText(
|
||||
sys.argv[2],
|
||||
sys.argv[3],
|
||||
)
|
||||
elif sys.argv[1] == "-t":
|
||||
GetTextTask(
|
||||
sys.argv[2],
|
||||
sys.argv[3],
|
||||
sys.argv[4],
|
||||
)
|
||||
else:
|
||||
print("Params: <runtime-config.json>")
|
||||
exit(0)
|
||||
@@ -1,50 +0,0 @@
|
||||
# -*- mode: python ; coding: utf-8 -*-
|
||||
|
||||
from PyInstaller.building.datastruct import Tree
|
||||
from PyInstaller.utils.hooks import get_package_paths
|
||||
|
||||
PACKAGE_DIRECTORY = get_package_paths('faster_whisper')[1]
|
||||
datas = [(PACKAGE_DIRECTORY, 'faster_whisper')]
|
||||
|
||||
|
||||
a = Analysis(
|
||||
['local_whisper.py'],
|
||||
pathex=[],
|
||||
binaries=[],
|
||||
datas=[],
|
||||
hiddenimports=[],
|
||||
hookspath=[],
|
||||
hooksconfig={},
|
||||
runtime_hooks=[],
|
||||
excludes=[],
|
||||
noarchive=False,
|
||||
optimize=0,
|
||||
)
|
||||
pyz = PYZ(a.pure)
|
||||
|
||||
exe = EXE(
|
||||
pyz,
|
||||
a.scripts,
|
||||
[],
|
||||
exclude_binaries=True,
|
||||
name='local_whisper',
|
||||
debug=False,
|
||||
bootloader_ignore_signals=False,
|
||||
strip=False,
|
||||
upx=True,
|
||||
console=True,
|
||||
disable_windowed_traceback=False,
|
||||
argv_emulation=False,
|
||||
target_arch=None,
|
||||
codesign_identity=None,
|
||||
entitlements_file=None,
|
||||
)
|
||||
coll = COLLECT(
|
||||
exe,
|
||||
a.binaries,
|
||||
a.datas,
|
||||
strip=False,
|
||||
upx=True,
|
||||
upx_exclude=[],
|
||||
name='local_whisper',
|
||||
)
|
||||
@@ -1,351 +0,0 @@
|
||||
# 读取文件的方法
|
||||
import json
|
||||
import os
|
||||
import win32api
|
||||
import win32con
|
||||
import pywintypes
|
||||
import shutil
|
||||
import re
|
||||
|
||||
|
||||
class PublicTools:
|
||||
"""
|
||||
一些公用的基础方法
|
||||
"""
|
||||
|
||||
def delete_path(self, path):
|
||||
"""
|
||||
删除指定路径的文件或者是文件夹
|
||||
"""
|
||||
# 检查路径是否存在
|
||||
if not os.path.exists(path):
|
||||
return
|
||||
|
||||
# 检查路径是文件还是文件夹
|
||||
if os.path.isfile(path):
|
||||
# 是文件，执行删除
|
||||
try:
|
||||
os.remove(path)
|
||||
except Exception as e:
|
||||
raise e
|
||||
elif os.path.isdir(path):
|
||||
# 是文件夹，执行删除
|
||||
try:
|
||||
shutil.rmtree(path)
|
||||
except Exception as e:
|
||||
raise e
|
||||
else:
|
||||
raise
|
||||
|
||||
def list_files_by_extension(self, folder_path, extension):
|
||||
"""
|
||||
读取指定文件夹下面的所有的指定拓展文件命的文件列表
|
||||
"""
|
||||
file_list = []
|
||||
for root, dirs, files in os.walk(folder_path):
|
||||
for file in files:
|
||||
if file.endswith(extension):
|
||||
file_list.append(os.path.join(root, file))
|
||||
elif file.endswith(extension.upper()):
|
||||
file_list.append(os.path.join(root, file))
|
||||
return file_list
|
||||
|
||||
def get_fonts_from_registry(self, key_path):
|
||||
"""
|
||||
获取注册表中安装的字体文件
|
||||
"""
|
||||
font_names = []
|
||||
try:
|
||||
key = win32api.RegOpenKeyEx(
|
||||
(
|
||||
win32con.HKEY_LOCAL_MACHINE
|
||||
if "HKEY_LOCAL_MACHINE" in key_path
|
||||
else win32con.HKEY_CURRENT_USER
|
||||
),
|
||||
key_path.split("\\", 1)[1],
|
||||
0,
|
||||
win32con.KEY_READ,
|
||||
)
|
||||
i = 0
|
||||
while True:
|
||||
try:
|
||||
value = win32api.RegEnumValue(key, i)
|
||||
font_name = value[0]
|
||||
# 使用正则表达式移除括号及其内容
|
||||
font_name = re.sub(r"\s*\([^)]*\)$", "", font_name)
|
||||
font_names.append(font_name)
|
||||
i += 1
|
||||
except pywintypes.error as e:
|
||||
if e.winerror == 259: # 没有更多的数据
|
||||
break
|
||||
else:
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
win32api.RegCloseKey(key)
|
||||
except:
|
||||
pass
|
||||
return font_names
|
||||
|
||||
def get_installed_fonts(self):
|
||||
"""
|
||||
获取字体文件名称并返回
|
||||
"""
|
||||
system_fonts = self.get_fonts_from_registry(
|
||||
"HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Fonts"
|
||||
)
|
||||
user_fonts = self.get_fonts_from_registry(
|
||||
"HKEY_CURRENT_USER\\Software\\Microsoft\\Windows NT\\CurrentVersion\\Fonts"
|
||||
)
|
||||
all_fonts = list(set(system_fonts + user_fonts)) # 合并并去重
|
||||
return all_fonts
|
||||
|
||||
# 将RRGGBB转换为BBGGRR
|
||||
def convert_rrggbb_to_bbggrr(self, rrggbb):
|
||||
"""
|
||||
将RRGGBB转换为BBGGRR
|
||||
"""
|
||||
if len(rrggbb) == 7:
|
||||
rr = rrggbb[1:3]
|
||||
gg = rrggbb[3:5]
|
||||
bb = rrggbb[5:7]
|
||||
return bb + gg + rr
|
||||
else:
|
||||
return "Invalid input"
|
||||
|
||||
def write_to_file(self, arr, filename):
|
||||
with open(filename, "w",encoding='utf-8') as f:
|
||||
for item in arr:
|
||||
f.write("%s\n" % item)
|
||||
|
||||
|
||||
# 读取文件
|
||||
def read_file(fileType):
|
||||
txt_path = input(f"输入{fileType}文件路径：")
|
||||
txt_path = remove_prefix_and_suffix(txt_path, '"', '"')
|
||||
while txt_path.strip() == "":
|
||||
txt_path = input(f"输入{fileType}文件路径：")
|
||||
|
||||
while os.path.exists(txt_path) == False:
|
||||
print("文件路径不存在错误：")
|
||||
txt_path = input(f"输入{fileType}文件路径：")
|
||||
txt_path = remove_prefix_and_suffix(txt_path, '"', '"')
|
||||
return txt_path
|
||||
|
||||
|
||||
def format_time_ms(milliseconds):
|
||||
"""
|
||||
时间转换将ms->小时:分钟:秒.毫秒格式
|
||||
"""
|
||||
seconds = milliseconds / 1000
|
||||
# 计算小时、分钟和秒
|
||||
hours = int(seconds // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
seconds = seconds % 60
|
||||
# 格式化字符串
|
||||
# 使用`%02d`确保小时和分钟总是显示为两位数，`%.2f`确保秒数显示两位小数
|
||||
formatted_time = f"{hours}:{minutes:02d}:{seconds:05.2f}"
|
||||
return formatted_time
|
||||
|
||||
|
||||
# 删除满足条件的开头和结尾
|
||||
def remove_prefix_and_suffix(input_str, prefix_to_remove, suffix_to_remove):
|
||||
if input_str.startswith(prefix_to_remove):
|
||||
# 删除开头
|
||||
input_str = input_str[len(prefix_to_remove) :]
|
||||
|
||||
if input_str.endswith(suffix_to_remove):
|
||||
# 删除结尾
|
||||
input_str = input_str[: -len(suffix_to_remove)]
|
||||
|
||||
return input_str
|
||||
|
||||
|
||||
# 判断文件夹下面是不是有特定的文件夹
|
||||
def check_if_folder_exists(parent_folder, target_folder_name):
|
||||
# 获取文件夹列表
|
||||
subfolders = [f.name for f in os.scandir(parent_folder) if f.is_dir()]
|
||||
|
||||
# 检查特定文件夹是否存在
|
||||
if target_folder_name in subfolders:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
# 检查指定文件夹中是否存在特定文件。
|
||||
def file_exists_in_folder(folder_path: str, file_name: str) -> bool:
|
||||
# 构建完整的文件路径
|
||||
file_path = os.path.join(folder_path, file_name)
|
||||
|
||||
# 返回文件是否存在
|
||||
return os.path.isfile(file_path)
|
||||
|
||||
|
||||
# 秒数转换，保留一位小数
|
||||
def convert_to_seconds(number, count):
|
||||
seconds = number / 1000000
|
||||
rounded_number = round(seconds, count)
|
||||
return rounded_number
|
||||
|
||||
|
||||
def is_empty(obj):
|
||||
if obj is None:
|
||||
return True
|
||||
elif isinstance(obj, str):
|
||||
return len(obj) == 0
|
||||
elif isinstance(obj, list):
|
||||
return len(obj) == 0
|
||||
elif isinstance(obj, dict):
|
||||
return len(obj) == 0
|
||||
return False
|
||||
|
||||
|
||||
def opt_dict(obj, key, default=None):
|
||||
if obj is None:
|
||||
return default
|
||||
if key in obj:
|
||||
v = obj[key]
|
||||
if not is_empty(v):
|
||||
return v
|
||||
return default
|
||||
|
||||
|
||||
def read_config(path, webui=True):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
runtime_config = json.load(f)
|
||||
|
||||
if "config" not in runtime_config:
|
||||
print("no filed 'config' in json")
|
||||
return None
|
||||
|
||||
config = runtime_config["config"]
|
||||
if "webui" not in config:
|
||||
print("no filed 'webui' in 'config'")
|
||||
return None
|
||||
|
||||
setting_config_path = config["setting"]
|
||||
if not os.path.exists(setting_config_path):
|
||||
setting_config_path = "config/" + setting_config_path
|
||||
if not os.path.exists(setting_config_path):
|
||||
setting_config_path = "../" + setting_config_path
|
||||
|
||||
# read config
|
||||
with open(setting_config_path, "r", encoding="utf-8") as f:
|
||||
setting_config = json.load(f)
|
||||
|
||||
# set workspace parent:根目录
|
||||
if "workspace" in setting_config:
|
||||
setting_config["workspace"]["parent"] = runtime_config["workspace"]
|
||||
else:
|
||||
setting_config["workspace"] = {"parent": runtime_config["workspace"]}
|
||||
setting_config["video"] = opt_dict(runtime_config, "video")
|
||||
|
||||
# merge setting config
|
||||
if "setting" in config:
|
||||
setting_config.update(runtime_config["setting"])
|
||||
|
||||
# webui config
|
||||
if webui:
|
||||
webui_config_path = config["webui"]
|
||||
if not os.path.exists(webui_config_path):
|
||||
webui_config_path = "config/webui/" + webui_config_path
|
||||
if not os.path.exists(webui_config_path):
|
||||
webui_config_path = "../" + webui_config_path
|
||||
|
||||
with open(webui_config_path, "r", encoding="utf-8") as f:
|
||||
webui_config = json.load(f)
|
||||
|
||||
# merge webui config
|
||||
if "webui" in runtime_config:
|
||||
webui_config.update(runtime_config["webui"])
|
||||
|
||||
return webui_config, setting_config
|
||||
return setting_config
|
||||
|
||||
|
||||
TAG_MODE_NONE = ""
|
||||
|
||||
|
||||
# 工作路径
|
||||
class Workspace:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: str,
|
||||
input: str,
|
||||
output: str,
|
||||
input_crop: str,
|
||||
output_crop: str,
|
||||
input_tag: str,
|
||||
input_mask: str,
|
||||
input_crop_mask: str,
|
||||
crop_info: str,
|
||||
):
|
||||
self.root = root
|
||||
self.input = input
|
||||
self.output = output
|
||||
self.input_crop = input_crop
|
||||
self.output_crop = output_crop
|
||||
self.input_tag = input_tag
|
||||
self.input_mask = input_mask
|
||||
self.input_crop_mask = input_crop_mask
|
||||
self.crop_info = crop_info
|
||||
|
||||
|
||||
# 定义一个倍数函数
|
||||
def round_up(num, mul):
|
||||
return (num // mul + 1) * mul
|
||||
|
||||
|
||||
class SettingConfig:
|
||||
|
||||
def __init__(self, config: dict, workParent):
|
||||
self.config = config
|
||||
self.webui_work_api = None
|
||||
self.workParent = workParent
|
||||
|
||||
def to_dict(self):
|
||||
return self.__dict__
|
||||
|
||||
def get_tag_mode(self):
|
||||
tag_cfg = opt_dict(self.config, "tag")
|
||||
return opt_dict(tag_cfg, "mode", TAG_MODE_NONE)
|
||||
|
||||
def get_tag_actions(self):
|
||||
tag_cfg = opt_dict(self.config, "tag")
|
||||
return opt_dict(tag_cfg, "actions", [])
|
||||
|
||||
def get_workspace_config(self) -> Workspace:
|
||||
workspace_config = opt_dict(self.config, "workspace")
|
||||
tmp_config = opt_dict(workspace_config, "tmp")
|
||||
|
||||
input = opt_dict(workspace_config, "input", "input")
|
||||
output = opt_dict(workspace_config, "output", "output")
|
||||
workspace_parent = self.workParent
|
||||
|
||||
tmp_parent = opt_dict(tmp_config, "parent", "tmp")
|
||||
input_crop = opt_dict(tmp_config, "input_crop", "input_crop")
|
||||
output_crop = opt_dict(tmp_config, "output_crop", "output_crop")
|
||||
input_tag = opt_dict(tmp_config, "input_tag", "input_crop")
|
||||
input_mask = opt_dict(tmp_config, "input_mask", "input_mask")
|
||||
input_crop_mask = opt_dict(tmp_config, "input_crop_mask", "input_crop_mask")
|
||||
crop_info = opt_dict(tmp_config, "crop_info", "crop_info.txt")
|
||||
|
||||
tmp_path = os.path.join(workspace_parent, tmp_parent)
|
||||
|
||||
return Workspace(
|
||||
workspace_parent,
|
||||
os.path.join(workspace_parent, input),
|
||||
os.path.join(workspace_parent, output),
|
||||
os.path.join(tmp_path, input_crop),
|
||||
os.path.join(tmp_path, output_crop),
|
||||
os.path.join(tmp_path, input_tag),
|
||||
os.path.join(tmp_path, input_mask),
|
||||
os.path.join(tmp_path, input_crop_mask),
|
||||
os.path.join(tmp_path, crop_info),
|
||||
)
|
||||
|
||||
def enable_tag(self):
|
||||
tag_cfg = opt_dict(self.config, "tag")
|
||||
return opt_dict(tag_cfg, "enable", True)
|
||||
@@ -1,307 +0,0 @@
|
||||
# pip install scenedetect opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple
|
||||
|
||||
from scenedetect.video_manager import VideoManager
|
||||
from scenedetect.scene_manager import SceneManager
|
||||
from scenedetect.stats_manager import StatsManager
|
||||
from scenedetect.detectors.content_detector import ContentDetector
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import subprocess
|
||||
from huggingface_hub import hf_hub_download
|
||||
from faster_whisper import WhisperModel
|
||||
from pathlib import Path
|
||||
import public_tools
|
||||
|
||||
# 获取智能画面分割的时间或者秒数
|
||||
def find_scenes(video_path, sensitivity):
|
||||
print(
|
||||
"正在计算分镜数据" + "sensitivity：" + str(sensitivity) + "path : " + video_path
|
||||
)
|
||||
sys.stdout.flush()
|
||||
video_manager = VideoManager([video_path])
|
||||
stats_manager = StatsManager()
|
||||
scene_manager = SceneManager(stats_manager)
|
||||
|
||||
# 使用contect-detector
|
||||
scene_manager.add_detector(ContentDetector(threshold=float(sensitivity)))
|
||||
|
||||
shijian_list = []
|
||||
|
||||
try:
|
||||
video_manager.set_downscale_factor()
|
||||
video_manager.start()
|
||||
scene_manager.detect_scenes(frame_source=video_manager)
|
||||
scene_list = scene_manager.get_scene_list()
|
||||
print("分镜数据列表：")
|
||||
sys.stdout.flush()
|
||||
for i, scene in enumerate(scene_list):
|
||||
shijian_list.append([scene[0].get_timecode(), scene[1].get_timecode()])
|
||||
print(
|
||||
"Scene %2d: Start %s / Frame %d, End %s / Frame %d"
|
||||
% (
|
||||
i + 1,
|
||||
scene[0].get_timecode(),
|
||||
scene[0].get_frames(),
|
||||
scene[1].get_timecode(),
|
||||
scene[1].get_frames(),
|
||||
)
|
||||
)
|
||||
sys.stdout.flush()
|
||||
finally:
|
||||
video_manager.release()
|
||||
|
||||
return shijian_list
|
||||
|
||||
|
||||
# 如果不存在就创建
|
||||
def createDir(file_dir):
|
||||
# 如果不存在文件夹，就创建
|
||||
if not os.path.isdir(file_dir):
|
||||
os.mkdir(file_dir)
|
||||
|
||||
|
||||
# 切分一个视频
|
||||
def ClipVideo(video_path, out_folder, image_out_folder, sensitivity, gpu_type):
|
||||
shijian_list = find_scenes(video_path, sensitivity) # 多组时间列表
|
||||
shijian_list_len = len(shijian_list)
|
||||
|
||||
print("总共有%s个场景" % str(shijian_list_len))
|
||||
sys.stdout.flush()
|
||||
video_list = []
|
||||
for i in range(0, shijian_list_len):
|
||||
start_time_str = shijian_list[i][0]
|
||||
end_time_str = shijian_list[i][1]
|
||||
|
||||
print("开始输出第" + str(i + 1) + "个分镜")
|
||||
video_name = "{:05d}".format(i + 1)
|
||||
out_video_file = os.path.join(out_folder, video_name + ".mp4")
|
||||
sys.stdout.flush()
|
||||
video_list.append(
|
||||
{
|
||||
"start_time_str": start_time_str,
|
||||
"end_time_str": end_time_str,
|
||||
"out_video_file": out_video_file,
|
||||
"video_name": video_name,
|
||||
}
|
||||
)
|
||||
|
||||
# 使用 ffmpeg 裁剪视频
|
||||
command = []
|
||||
command.append("ffmpeg")
|
||||
command.append("-i")
|
||||
command.append(video_path)
|
||||
command.append("-ss")
|
||||
command.append(start_time_str)
|
||||
command.append("-to")
|
||||
command.append(end_time_str)
|
||||
command.append("-c:v")
|
||||
|
||||
if gpu_type == "NVIDIA":
|
||||
command.append("h264_nvenc")
|
||||
elif gpu_type == "AMD":
|
||||
command.append("h264_amf")
|
||||
else:
|
||||
command.append("libx264")
|
||||
|
||||
command.append("-preset")
|
||||
command.append("fast")
|
||||
command.append("-c:a")
|
||||
command.append("copy")
|
||||
command.append(out_video_file)
|
||||
command.append("-loglevel")
|
||||
command.append("error")
|
||||
|
||||
subprocess.run(
|
||||
command,
|
||||
check=True,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
|
||||
print("分镜输出完成。开始抽帧")
|
||||
sys.stdout.flush()
|
||||
for vi in video_list:
|
||||
h, m, s = vi["start_time_str"].split(":")
|
||||
start_seconds = int(h) * 3600 + int(m) * 60 + float(s)
|
||||
|
||||
h, m, s = vi["end_time_str"].split(":")
|
||||
end_seconds = int(h) * 3600 + int(m) * 60 + float(s)
|
||||
print("正在抽帧：" + vi["video_name"])
|
||||
sys.stdout.flush()
|
||||
subprocess.run(
|
||||
[
|
||||
"ffmpeg",
|
||||
"-ss",
|
||||
str((end_seconds - start_seconds) / 2),
|
||||
"-i",
|
||||
vi["out_video_file"],
|
||||
"-frames:v",
|
||||
"1",
|
||||
os.path.join(image_out_folder, vi["video_name"] + ".png"),
|
||||
"-loglevel",
|
||||
"error",
|
||||
]
|
||||
)
|
||||
|
||||
print("抽帧完成，开始识别文案")
|
||||
sys.stdout.flush()
|
||||
return video_list
|
||||
|
||||
|
||||
def SplitAudio(video_out_folder, video_list):
|
||||
# ffmpeg -i input_file.mp4 -vn -ab 128k output_file.mp3
|
||||
print("正在分离音频!!")
|
||||
mp3_list = []
|
||||
sys.stdout.flush()
|
||||
for v in video_list:
|
||||
mp3_path = os.path.join(video_out_folder, v["video_name"] + ".mp3")
|
||||
mp3_list.append(mp3_path)
|
||||
subprocess.run(
|
||||
[
|
||||
"ffmpeg",
|
||||
"-i",
|
||||
v["out_video_file"],
|
||||
"-vn",
|
||||
"-ab",
|
||||
"128k",
|
||||
mp3_path,
|
||||
"-loglevel",
|
||||
"error",
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
return mp3_list
|
||||
|
||||
|
||||
def GetText(out_folder, mp3_list):
|
||||
text = []
|
||||
# 先获取模型
|
||||
print("正在下载或加载模型")
|
||||
sys.stdout.flush()
|
||||
model_path = Path(
|
||||
hf_hub_download(repo_id="Systran/faster-whisper-large-v3", filename="model.bin")
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="preprocessor_config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="tokenizer.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="vocabulary.json",
|
||||
)
|
||||
model = WhisperModel(
|
||||
model_size_or_path=os.path.dirname(model_path),
|
||||
device="auto",
|
||||
local_files_only=True,
|
||||
)
|
||||
print("模型加载成功，开始识别")
|
||||
sys.stdout.flush()
|
||||
for mp in mp3_list:
|
||||
segments, info = model.transcribe(
|
||||
mp,
|
||||
beam_size=5,
|
||||
language="zh",
|
||||
vad_filter=True,
|
||||
vad_parameters=dict(min_silence_duration_ms=1000),
|
||||
)
|
||||
tmp_text = ""
|
||||
for segment in segments:
|
||||
tmp_text += segment.text + "。"
|
||||
print(mp + "识别完成")
|
||||
sys.stdout.flush()
|
||||
text.append(tmp_text)
|
||||
|
||||
# 数据写出
|
||||
print("文本全部识别成功，正在写出")
|
||||
sys.stdout.flush()
|
||||
tools = public_tools.PublicTools()
|
||||
tools.write_to_file(text, os.path.join(out_folder, "文案.txt"))
|
||||
print("写出完成")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def GetTextTask(out_folder, mp, name):
|
||||
text = []
|
||||
# 先获取模型
|
||||
print("正在下载或加载模型")
|
||||
sys.stdout.flush()
|
||||
model_path = Path(
|
||||
hf_hub_download(repo_id="Systran/faster-whisper-large-v3", filename="model.bin")
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="preprocessor_config.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="tokenizer.json",
|
||||
)
|
||||
hf_hub_download(
|
||||
repo_id="Systran/faster-whisper-large-v3",
|
||||
filename="vocabulary.json",
|
||||
)
|
||||
model = WhisperModel(
|
||||
model_size_or_path=os.path.dirname(model_path),
|
||||
device="auto",
|
||||
local_files_only=True,
|
||||
)
|
||||
print("模型加载成功，开始识别")
|
||||
sys.stdout.flush()
|
||||
segments, info = model.transcribe(
|
||||
mp,
|
||||
beam_size=5,
|
||||
language="zh",
|
||||
vad_filter=True,
|
||||
vad_parameters=dict(min_silence_duration_ms=1000),
|
||||
)
|
||||
tmp_text = ""
|
||||
for segment in segments:
|
||||
tmp_text += segment.text + "。"
|
||||
print(mp + "识别完成")
|
||||
sys.stdout.flush()
|
||||
text.append(tmp_text)
|
||||
|
||||
# 数据写出
|
||||
sys.stdout.flush()
|
||||
tools = public_tools.PublicTools()
|
||||
tools.write_to_file(text, os.path.join(out_folder, name + ".txt"))
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def get_fram(video_path, out_path, sensitivity):
|
||||
try:
|
||||
shijian_list = find_scenes(video_path, sensitivity) # 多组时间列表
|
||||
print("总共有%s个场景" % str(len(shijian_list)))
|
||||
print("开始输出json")
|
||||
print(shijian_list)
|
||||
# 将数组中的消息写道json文件中
|
||||
with open(out_path, "w") as file:
|
||||
# 将数组写入到指定的json文件
|
||||
json.dump(shijian_list, file)
|
||||
print("输出完成")
|
||||
except Exception as e:
|
||||
print("出现错误" + str(e))
|
||||
exit(0)
|
||||
|
||||
|
||||
def init(video_path, video_out_folder, image_out_folder, sensitivity, gpu_type):
|
||||
v_l = ClipVideo(
|
||||
video_path, video_out_folder, image_out_folder, sensitivity, gpu_type
|
||||
)
|
||||
|
||||
# 开始分离音频
|
||||
m_l = SplitAudio(video_out_folder, v_l)
|
||||
# 开始识别字幕
|
||||
GetText(os.path.dirname(video_out_folder), m_l)
|
||||
Reference in New Issue
Block a user