101 lines
4.2 KiB
Python
101 lines
4.2 KiB
Python
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import transformers
|
|
from diffusers import StableDiffusionDepth2ImgPipeline, StableDiffusionXLPipeline
|
|
from transformers import Owlv2Processor, Owlv2ForObjectDetection
|
|
|
|
from data_generation import words_bank
|
|
|
|
|
|
def main():
|
|
sdxl_pipe = StableDiffusionXLPipeline.from_pretrained(
|
|
"stabilityai/sdxl-turbo", use_safetensors=True
|
|
).to("cuda")
|
|
|
|
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-2-depth", torch_dtype=torch.float16, ).to("cuda")
|
|
|
|
processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
|
|
model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble")
|
|
|
|
bria_pipe = transformers.pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
|
|
|
|
out_dir = Path(f'datasets/generated/texture_data')
|
|
out_dir.mkdir(exist_ok=True, parents=True)
|
|
|
|
with open('assets/openimages_classes.txt', 'r') as f:
|
|
objects = f.read().splitlines()
|
|
|
|
for _ in range(100000):
|
|
object_name = random.choice(objects)
|
|
# Remove special characters
|
|
object_name = ''.join(char if char.isalnum() else ' ' for char in object_name)
|
|
# Restrict to two words
|
|
object_name = ' '.join(object_name.split()[:2])
|
|
placement = random.choice(words_bank.placements) if random.random() < 0.5 else ''
|
|
prompt = f"A {object_name} {placement}"
|
|
|
|
seed = random.randint(0, 1000000)
|
|
object_out_dir = out_dir / f'{object_name.replace(" ", "_")}_{placement.replace(" ", "_")}_{seed}'
|
|
if object_out_dir.exists():
|
|
continue
|
|
object_out_dir.mkdir(exist_ok=True, parents=True)
|
|
|
|
base_image = sdxl_pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
|
|
out_path = object_out_dir / f'base.jpg'
|
|
base_image.save(out_path)
|
|
|
|
# Find box
|
|
texts = [[f"a {object_name}"]]
|
|
inputs = processor(text=texts, images=base_image, return_tensors="pt")
|
|
outputs = model(**inputs)
|
|
|
|
# Target image sizes (height, width) to rescale box predictions [batch_size, 2]
|
|
target_sizes = torch.Tensor([base_image.size[::-1]])
|
|
# Convert outputs (bounding boxes and class logits) to Pascal VOC Format (xmin, ymin, xmax, ymax)
|
|
results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes, threshold=0.1)
|
|
|
|
i = 0 # Retrieve predictions for the first image for the corresponding text queries
|
|
boxes, scores, labels = results[i]["boxes"], results[i]["scores"], results[i]["labels"]
|
|
# Filter to box with max score
|
|
if len(boxes) == 0:
|
|
scores = torch.tensor([1.0])
|
|
boxes = torch.tensor([[0, 0, base_image.size[0], base_image.size[1]]])
|
|
|
|
max_score_idx = scores.argmax()
|
|
box = boxes[max_score_idx]
|
|
|
|
# Save box to json file
|
|
box_dict = {
|
|
"x1": int(box[0].item()),
|
|
"y1": int(box[1].item()),
|
|
"x2": int(box[2].item()),
|
|
"y2": int(box[3].item()),
|
|
}
|
|
crop = base_image.crop((box_dict['x1'], box_dict['y1'], box_dict['x2'], box_dict['y2']))
|
|
tmp_path = 'tmp.jpg'
|
|
crop.save(tmp_path)
|
|
crop_mask = bria_pipe(tmp_path, return_mask=True)
|
|
crop_mask.save(object_out_dir / 'mask.png')
|
|
with open(object_out_dir / 'box.json', 'w') as f:
|
|
json.dump(box_dict, f)
|
|
|
|
for _ in range(5):
|
|
num_samples = random.randint(1, 5)
|
|
sample_attributes = random.sample(words_bank.texture_attributes, num_samples)
|
|
prompt = f"A {object_name} made from {' '.join(sample_attributes)} {placement}"
|
|
n_propmt = "bad, deformed, ugly, bad anotomy"
|
|
seed = random.randint(0, 1000000)
|
|
image = pipe(prompt=prompt, image=base_image, negative_prompt=n_propmt, strength=1.0,
|
|
generator=torch.Generator().manual_seed(seed)).images[0]
|
|
attrs_str = '_'.join(sample_attributes)[:100]
|
|
out_path = object_out_dir / f'{attrs_str}_{seed}.jpg'
|
|
image.save(out_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|