Added dataset tagging and management tools using llava
This commit is contained in:
Jaret Burkett
29
extensions_built_in/dataset_tools/tools/image_tools.py
Normal file
29
extensions_built_in/dataset_tools/tools/image_tools.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from typing import Literal, Type
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
Step: Type = Literal['caption', 'caption_short', 'create_mask', 'contrast_stretch']
|
||||
|
||||
img_manipulation_steps = ['contrast_stretch']
|
||||
|
||||
|
||||
def pil_to_cv2(image):
|
||||
"""Convert a PIL image to a cv2 image."""
|
||||
return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
||||
|
||||
|
||||
def cv2_to_pil(image):
|
||||
"""Convert a cv2 image to a PIL image."""
|
||||
return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
def load_image(img_path: str):
|
||||
image = Image.open(img_path).convert('RGB')
|
||||
try:
|
||||
# transpose with exif data
|
||||
image = ImageOps.exif_transpose(image)
|
||||
except Exception as e:
|
||||
pass
|
||||
return image
|
||||
|
||||
129
extensions_built_in/dataset_tools/tools/llava_utils.py
Normal file
129
extensions_built_in/dataset_tools/tools/llava_utils.py
Normal file
@@ -0,0 +1,129 @@
|
||||
try:
|
||||
from llava.model import LlavaLlamaForCausalLM
|
||||
except ImportError:
|
||||
# print("You need to manually install llava -> pip install --no-deps git+https://github.com/haotian-liu/LLaVA.git")
|
||||
print("You need to manually install llava -> pip install --no-deps git+https://github.com/haotian-liu/LLaVA.git")
|
||||
raise
|
||||
|
||||
long_prompt = 'caption this image. describe every single thing in the image in detail. Do not include any unnecessary words in your description for the sake of good grammar. I want many short statements that serve the single purpose of giving the most thorough description if items as possible in the smallest, comma separated way possible. be sure to describe people\'s moods, clothing, the environment, lighting, colors, and everything.'
|
||||
short_prompt = 'caption this image in less than ten words'
|
||||
|
||||
prompts = [
|
||||
long_prompt,
|
||||
short_prompt,
|
||||
]
|
||||
|
||||
replacements = [
|
||||
("the image features", ""),
|
||||
("the image shows", ""),
|
||||
("the image depicts", ""),
|
||||
("the image is", ""),
|
||||
]
|
||||
|
||||
import torch
|
||||
from PIL import Image, ImageOps
|
||||
from llava.conversation import conv_templates, SeparatorStyle
|
||||
from transformers import AutoTokenizer, BitsAndBytesConfig, CLIPImageProcessor
|
||||
from llava.utils import disable_torch_init
|
||||
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
|
||||
from llava.mm_utils import tokenizer_image_token, KeywordsStoppingCriteria
|
||||
|
||||
img_ext = ['.jpg', '.jpeg', '.png', '.webp']
|
||||
|
||||
|
||||
class LLaVAImageProcessor:
|
||||
def __init__(self, device='cuda'):
|
||||
self.device = device
|
||||
self.model: LlavaLlamaForCausalLM = None
|
||||
self.tokenizer: AutoTokenizer = None
|
||||
self.image_processor: CLIPImageProcessor = None
|
||||
self.is_loaded = False
|
||||
|
||||
def load_model(self):
|
||||
from llava.model import LlavaLlamaForCausalLM
|
||||
|
||||
model_path = "4bit/llava-v1.5-13b-3GB"
|
||||
# kwargs = {"device_map": "auto"}
|
||||
kwargs = {"device_map": self.device}
|
||||
kwargs['load_in_4bit'] = True
|
||||
kwargs['quantization_config'] = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
bnb_4bit_quant_type='nf4'
|
||||
)
|
||||
self.model = LlavaLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
|
||||
vision_tower = self.model.get_vision_tower()
|
||||
if not vision_tower.is_loaded:
|
||||
vision_tower.load_model()
|
||||
vision_tower.to(device=self.device)
|
||||
self.image_processor = vision_tower.image_processor
|
||||
self.is_loaded = True
|
||||
|
||||
def clean_caption(self, cap):
|
||||
# remove any newlines
|
||||
cap = cap.replace("\n", ", ")
|
||||
cap = cap.replace("\r", ", ")
|
||||
cap = cap.replace(".", ",")
|
||||
cap = cap.replace("\"", "")
|
||||
|
||||
# remove unicode characters
|
||||
cap = cap.encode('ascii', 'ignore').decode('ascii')
|
||||
|
||||
# make lowercase
|
||||
cap = cap.lower()
|
||||
# remove any extra spaces
|
||||
cap = " ".join(cap.split())
|
||||
|
||||
for replacement in replacements:
|
||||
cap = cap.replace(replacement[0], replacement[1])
|
||||
|
||||
cap_list = cap.split(",")
|
||||
# trim whitespace
|
||||
cap_list = [c.strip() for c in cap_list]
|
||||
# remove empty strings
|
||||
cap_list = [c for c in cap_list if c != ""]
|
||||
# remove duplicates
|
||||
cap_list = list(dict.fromkeys(cap_list))
|
||||
# join back together
|
||||
cap = ", ".join(cap_list)
|
||||
return cap
|
||||
|
||||
def generate_caption(self, image: Image, prompt: str = long_prompt):
|
||||
# question = "how many dogs are in the picture?"
|
||||
disable_torch_init()
|
||||
conv_mode = "llava_v0"
|
||||
conv = conv_templates[conv_mode].copy()
|
||||
roles = conv.roles
|
||||
image_tensor = self.image_processor.preprocess([image], return_tensors='pt')['pixel_values'].half().cuda()
|
||||
|
||||
inp = f"{roles[0]}: {prompt}"
|
||||
inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp
|
||||
conv.append_message(conv.roles[0], inp)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
raw_prompt = conv.get_prompt()
|
||||
input_ids = tokenizer_image_token(raw_prompt, self.tokenizer, IMAGE_TOKEN_INDEX,
|
||||
return_tensors='pt').unsqueeze(0).cuda()
|
||||
stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
|
||||
keywords = [stop_str]
|
||||
stopping_criteria = KeywordsStoppingCriteria(keywords, self.tokenizer, input_ids)
|
||||
with torch.inference_mode():
|
||||
output_ids = self.model.generate(
|
||||
input_ids, images=image_tensor, do_sample=True, temperature=0.1,
|
||||
max_new_tokens=1024, use_cache=True, stopping_criteria=[stopping_criteria],
|
||||
top_p=0.9
|
||||
)
|
||||
outputs = self.tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
|
||||
conv.messages[-1][-1] = outputs
|
||||
output = outputs.rsplit('</s>', 1)[0]
|
||||
return self.clean_caption(output)
|
||||
|
||||
def generate_captions(self, image: Image):
|
||||
|
||||
responses = []
|
||||
for prompt in prompts:
|
||||
output = self.generate_caption(image, prompt)
|
||||
responses.append(output)
|
||||
# replace all . with ,
|
||||
return responses
|
||||
279
extensions_built_in/dataset_tools/tools/sync_tools.py
Normal file
279
extensions_built_in/dataset_tools/tools/sync_tools.py
Normal file
@@ -0,0 +1,279 @@
|
||||
import os
|
||||
import requests
|
||||
import tqdm
|
||||
from typing import List, Optional, TYPE_CHECKING
|
||||
|
||||
|
||||
def img_root_path(img_id: str):
|
||||
return os.path.dirname(os.path.dirname(img_id))
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..dataset_tools_config_modules import DatasetSyncCollectionConfig
|
||||
|
||||
img_exts = ['.jpg', '.jpeg', '.webp', '.png']
|
||||
|
||||
class Photo:
|
||||
def __init__(
|
||||
self,
|
||||
id,
|
||||
host,
|
||||
width,
|
||||
height,
|
||||
url,
|
||||
filename
|
||||
):
|
||||
self.id = str(id)
|
||||
self.host = host
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.url = url
|
||||
self.filename = filename
|
||||
|
||||
|
||||
def get_desired_size(img_width: int, img_height: int, min_width: int, min_height: int):
|
||||
if img_width > img_height:
|
||||
scale = min_height / img_height
|
||||
else:
|
||||
scale = min_width / img_width
|
||||
|
||||
new_width = int(img_width * scale)
|
||||
new_height = int(img_height * scale)
|
||||
|
||||
return new_width, new_height
|
||||
|
||||
|
||||
def get_pexels_images(config: 'DatasetSyncCollectionConfig') -> List[Photo]:
|
||||
all_images = []
|
||||
next_page = f"https://api.pexels.com/v1/collections/{config.collection_id}?page=1&per_page=80&type=photos"
|
||||
|
||||
while True:
|
||||
response = requests.get(next_page, headers={
|
||||
"Authorization": f"{config.api_key}"
|
||||
})
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
all_images.extend(data['media'])
|
||||
if 'next_page' in data and data['next_page']:
|
||||
next_page = data['next_page']
|
||||
else:
|
||||
break
|
||||
|
||||
photos = []
|
||||
for image in all_images:
|
||||
new_width, new_height = get_desired_size(image['width'], image['height'], config.min_width, config.min_height)
|
||||
url = f"{image['src']['original']}?auto=compress&cs=tinysrgb&h={new_height}&w={new_width}"
|
||||
filename = os.path.basename(image['src']['original'])
|
||||
|
||||
photos.append(Photo(
|
||||
id=image['id'],
|
||||
host="pexels",
|
||||
width=image['width'],
|
||||
height=image['height'],
|
||||
url=url,
|
||||
filename=filename
|
||||
))
|
||||
|
||||
return photos
|
||||
|
||||
|
||||
def get_unsplash_images(config: 'DatasetSyncCollectionConfig') -> List[Photo]:
|
||||
headers = {
|
||||
# "Authorization": f"Client-ID {UNSPLASH_ACCESS_KEY}"
|
||||
"Authorization": f"Client-ID {config.api_key}"
|
||||
}
|
||||
# headers['Authorization'] = f"Bearer {token}"
|
||||
|
||||
url = f"https://api.unsplash.com/collections/{config.collection_id}/photos?page=1&per_page=30"
|
||||
response = requests.get(url, headers=headers)
|
||||
response.raise_for_status()
|
||||
res_headers = response.headers
|
||||
# parse the link header to get the next page
|
||||
# 'Link': '<https://api.unsplash.com/collections/mIPWwLdfct8/photos?page=82>; rel="last", <https://api.unsplash.com/collections/mIPWwLdfct8/photos?page=2>; rel="next"'
|
||||
has_next_page = False
|
||||
if 'Link' in res_headers:
|
||||
has_next_page = True
|
||||
link_header = res_headers['Link']
|
||||
link_header = link_header.split(',')
|
||||
link_header = [link.strip() for link in link_header]
|
||||
link_header = [link.split(';') for link in link_header]
|
||||
link_header = [[link[0].strip('<>'), link[1].strip().strip('"')] for link in link_header]
|
||||
link_header = {link[1]: link[0] for link in link_header}
|
||||
|
||||
# get page number from last url
|
||||
last_page = link_header['rel="last']
|
||||
last_page = last_page.split('?')[1]
|
||||
last_page = last_page.split('&')
|
||||
last_page = [param.split('=') for param in last_page]
|
||||
last_page = {param[0]: param[1] for param in last_page}
|
||||
last_page = int(last_page['page'])
|
||||
|
||||
all_images = response.json()
|
||||
|
||||
if has_next_page:
|
||||
# assume we start on page 1, so we don't need to get it again
|
||||
for page in tqdm.tqdm(range(2, last_page + 1)):
|
||||
url = f"https://api.unsplash.com/collections/{config.collection_id}/photos?page={page}&per_page=30"
|
||||
response = requests.get(url, headers=headers)
|
||||
response.raise_for_status()
|
||||
all_images.extend(response.json())
|
||||
|
||||
photos = []
|
||||
for image in all_images:
|
||||
new_width, new_height = get_desired_size(image['width'], image['height'], config.min_width, config.min_height)
|
||||
url = f"{image['urls']['raw']}&w={new_width}"
|
||||
filename = f"{image['id']}.jpg"
|
||||
|
||||
photos.append(Photo(
|
||||
id=image['id'],
|
||||
host="unsplash",
|
||||
width=image['width'],
|
||||
height=image['height'],
|
||||
url=url,
|
||||
filename=filename
|
||||
))
|
||||
|
||||
return photos
|
||||
|
||||
|
||||
def get_img_paths(dir_path: str):
|
||||
os.makedirs(dir_path, exist_ok=True)
|
||||
local_files = os.listdir(dir_path)
|
||||
# remove non image files
|
||||
local_files = [file for file in local_files if os.path.splitext(file)[1].lower() in img_exts]
|
||||
# make full path
|
||||
local_files = [os.path.join(dir_path, file) for file in local_files]
|
||||
return local_files
|
||||
|
||||
|
||||
def get_local_image_ids(dir_path: str):
|
||||
os.makedirs(dir_path, exist_ok=True)
|
||||
local_files = get_img_paths(dir_path)
|
||||
# assuming local files are named after Unsplash IDs, e.g., 'abc123.jpg'
|
||||
return set([os.path.basename(file).split('.')[0] for file in local_files])
|
||||
|
||||
|
||||
def get_local_image_file_names(dir_path: str):
|
||||
os.makedirs(dir_path, exist_ok=True)
|
||||
local_files = get_img_paths(dir_path)
|
||||
# assuming local files are named after Unsplash IDs, e.g., 'abc123.jpg'
|
||||
return set([os.path.basename(file) for file in local_files])
|
||||
|
||||
|
||||
def download_image(photo: Photo, dir_path: str, min_width: int = 1024, min_height: int = 1024):
|
||||
img_width = photo.width
|
||||
img_height = photo.height
|
||||
|
||||
if img_width < min_width or img_height < min_height:
|
||||
raise ValueError(f"Skipping {photo.id} because it is too small: {img_width}x{img_height}")
|
||||
|
||||
img_response = requests.get(photo.url)
|
||||
img_response.raise_for_status()
|
||||
os.makedirs(dir_path, exist_ok=True)
|
||||
|
||||
filename = os.path.join(dir_path, photo.filename)
|
||||
with open(filename, 'wb') as file:
|
||||
file.write(img_response.content)
|
||||
|
||||
|
||||
def update_caption(img_path: str):
|
||||
# if the caption is a txt file, convert it to a json file
|
||||
filename_no_ext = os.path.splitext(os.path.basename(img_path))[0]
|
||||
# see if it exists
|
||||
if os.path.exists(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.json")):
|
||||
# todo add poi and what not
|
||||
return # we have a json file
|
||||
caption = ""
|
||||
# see if txt file exists
|
||||
if os.path.exists(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.txt")):
|
||||
# read it
|
||||
with open(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.txt"), 'r') as file:
|
||||
caption = file.read()
|
||||
# write json file
|
||||
with open(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.json"), 'w') as file:
|
||||
file.write(f'{{"caption": "{caption}"}}')
|
||||
|
||||
# delete txt file
|
||||
os.remove(os.path.join(os.path.dirname(img_path), f"{filename_no_ext}.txt"))
|
||||
|
||||
|
||||
# def equalize_img(img_path: str):
|
||||
# input_path = img_path
|
||||
# output_path = os.path.join(img_root_path(img_path), COLOR_CORRECTED_DIR, os.path.basename(img_path))
|
||||
# os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
# process_img(
|
||||
# img_path=input_path,
|
||||
# output_path=output_path,
|
||||
# equalize=True,
|
||||
# max_size=2056,
|
||||
# white_balance=False,
|
||||
# gamma_correction=False,
|
||||
# strength=0.6,
|
||||
# )
|
||||
|
||||
|
||||
# def annotate_depth(img_path: str):
|
||||
# # make fake args
|
||||
# args = argparse.Namespace()
|
||||
# args.annotator = "midas"
|
||||
# args.res = 1024
|
||||
#
|
||||
# img = cv2.imread(img_path)
|
||||
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
#
|
||||
# output = annotate(img, args)
|
||||
#
|
||||
# output = output.astype('uint8')
|
||||
# output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
|
||||
#
|
||||
# os.makedirs(os.path.dirname(img_path), exist_ok=True)
|
||||
# output_path = os.path.join(img_root_path(img_path), DEPTH_DIR, os.path.basename(img_path))
|
||||
#
|
||||
# cv2.imwrite(output_path, output)
|
||||
|
||||
|
||||
# def invert_depth(img_path: str):
|
||||
# img = cv2.imread(img_path)
|
||||
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
# # invert the colors
|
||||
# img = cv2.bitwise_not(img)
|
||||
#
|
||||
# os.makedirs(os.path.dirname(img_path), exist_ok=True)
|
||||
# output_path = os.path.join(img_root_path(img_path), INVERTED_DEPTH_DIR, os.path.basename(img_path))
|
||||
# cv2.imwrite(output_path, img)
|
||||
|
||||
|
||||
#
|
||||
# # update our list of raw images
|
||||
# raw_images = get_img_paths(raw_dir)
|
||||
#
|
||||
# # update raw captions
|
||||
# for image_id in tqdm.tqdm(raw_images, desc="Updating raw captions"):
|
||||
# update_caption(image_id)
|
||||
#
|
||||
# # equalize images
|
||||
# for img_path in tqdm.tqdm(raw_images, desc="Equalizing images"):
|
||||
# if img_path not in eq_images:
|
||||
# equalize_img(img_path)
|
||||
#
|
||||
# # update our list of eq images
|
||||
# eq_images = get_img_paths(eq_dir)
|
||||
# # update eq captions
|
||||
# for image_id in tqdm.tqdm(eq_images, desc="Updating eq captions"):
|
||||
# update_caption(image_id)
|
||||
#
|
||||
# # annotate depth
|
||||
# depth_dir = os.path.join(root_dir, DEPTH_DIR)
|
||||
# depth_images = get_img_paths(depth_dir)
|
||||
# for img_path in tqdm.tqdm(eq_images, desc="Annotating depth"):
|
||||
# if img_path not in depth_images:
|
||||
# annotate_depth(img_path)
|
||||
#
|
||||
# depth_images = get_img_paths(depth_dir)
|
||||
#
|
||||
# # invert depth
|
||||
# inv_depth_dir = os.path.join(root_dir, INVERTED_DEPTH_DIR)
|
||||
# inv_depth_images = get_img_paths(inv_depth_dir)
|
||||
# for img_path in tqdm.tqdm(depth_images, desc="Inverting depth"):
|
||||
# if img_path not in inv_depth_images:
|
||||
# invert_depth(img_path)
|
||||
Reference in New Issue
Block a user