160 lines
7.0 KiB
Python
160 lines
7.0 KiB
Python
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import transformers
|
|
from diffusers import 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")
|
|
|
|
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/union_data')
|
|
out_dir.mkdir(exist_ok=True, parents=True)
|
|
|
|
# Generated images that did not meet the criteria
|
|
out_dir_leftovers = Path(f'datasets/generated/union_data_leftovers')
|
|
out_dir_leftovers.mkdir(exist_ok=True, parents=True)
|
|
|
|
with open('assets/openimages_classes.txt', 'r') as f:
|
|
objects = f.read().splitlines()
|
|
|
|
base_image = None
|
|
saved_base = False
|
|
for _ in range(100000):
|
|
if base_image is not None and not saved_base:
|
|
alternative_path = out_dir_leftovers / f'{object_name.replace(" ", "_")}_{placement.replace(" ", "_")}_{seed}.jpg'
|
|
base_image.save(alternative_path)
|
|
saved_base = False
|
|
try:
|
|
current_objects = []
|
|
object_count = 2
|
|
for _ in range(object_count):
|
|
object_name = random.choice(objects)
|
|
# Remove '/' or any non english 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])
|
|
current_objects.append(object_name)
|
|
object_name = ' and a '.join(current_objects)
|
|
print(object_name)
|
|
placement = random.choice(words_bank.placements) if random.random() < 0.5 else ''
|
|
|
|
prompt = f"A photo of {object_name} {placement}"
|
|
seed = random.randint(0, 1000000)
|
|
|
|
base_image = sdxl_pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
|
|
object_name = object_name[:50]
|
|
out_path = out_dir / f'{object_name.replace(" ", "_")}_{placement.replace(" ", "_")}_{seed}.jpg'
|
|
|
|
# Try to detect the objects in the generated image
|
|
texts = [[f"a {obj}" for obj in current_objects]]
|
|
inputs = processor(text=texts, images=base_image, return_tensors="pt")
|
|
with torch.no_grad():
|
|
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.2)
|
|
|
|
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"]
|
|
|
|
chosen_boxes = []
|
|
# Take highest scoring box for each label
|
|
for obj_ind in range(len(current_objects)):
|
|
relevant_boxes = boxes[labels == obj_ind]
|
|
relevant_scores = scores[labels == obj_ind]
|
|
if len(relevant_boxes) > 0:
|
|
max_score_idx = relevant_scores.argmax()
|
|
max_box = relevant_boxes[max_score_idx]
|
|
if relevant_scores[max_score_idx] < 0.2:
|
|
break
|
|
chosen_boxes.append(max_box)
|
|
|
|
# Filter to box with max score
|
|
if len(chosen_boxes) != 2:
|
|
print(f'Skipping, detected {len(chosen_boxes)} objects')
|
|
continue
|
|
|
|
# Verify small overlap between two using IoU
|
|
box1, box2 = chosen_boxes
|
|
box1 = box1.int()
|
|
box2 = box2.int()
|
|
x1 = max(box1[0], box2[0])
|
|
y1 = max(box1[1], box2[1])
|
|
x2 = min(box1[2], box2[2])
|
|
y2 = min(box1[3], box2[3])
|
|
intersection = max(0, x2 - x1) * max(0, y2 - y1)
|
|
area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
|
|
area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
|
|
union = area1 + area2 - intersection
|
|
iou = float(intersection / union)
|
|
if iou > 0.1:
|
|
print(f'Skipping, iou is {iou}')
|
|
continue
|
|
|
|
metadata = {'objects': [], 'placement': placement, 'prompt': prompt}
|
|
masked_objects = []
|
|
for obj_ind, box in enumerate(chosen_boxes):
|
|
box = box.int()
|
|
obj_name = current_objects[obj_ind]
|
|
box_dict = {
|
|
"x1": int(box[0].item()),
|
|
"y1": int(box[1].item()),
|
|
"x2": int(box[2].item()),
|
|
"y2": int(box[3].item()),
|
|
"name": obj_name
|
|
}
|
|
object_size = (box_dict['x2'] - box_dict['x1']) * (box_dict['y2'] - box_dict['y1'])
|
|
if object_size < 5000:
|
|
print(f'Skipping, object size is {object_size}')
|
|
continue
|
|
metadata['objects'].append(box_dict)
|
|
crop = base_image.crop((box_dict['x1'], box_dict['y1'], box_dict['x2'], box_dict['y2']))
|
|
tmp_path = 'tmp.jpg'
|
|
crop.save(tmp_path)
|
|
with torch.no_grad():
|
|
masked_object = bria_pipe(tmp_path)
|
|
# Returns RGBA, take mask channel and check how many pixels
|
|
crop_mask = np.array(masked_object)[..., 3]
|
|
|
|
# Make sure at least 10000 pixels are non-masked
|
|
# Check if over half of the pixels seen
|
|
total_pixels = crop_mask.size
|
|
seen_pixels = np.sum(crop_mask > 200)
|
|
if seen_pixels / total_pixels < 0.1:
|
|
print(f'Skipping, not enough pixels seen. only {seen_pixels / total_pixels:.3f} seen')
|
|
continue
|
|
masked_objects.append(masked_object)
|
|
if len(masked_objects) != len(chosen_boxes):
|
|
continue
|
|
for obj_ind, masked_object in enumerate(masked_objects):
|
|
masked_object.save(str(out_path).replace('.jpg', f'_OBJ_{obj_ind}.png'))
|
|
with open(out_path.with_suffix('.json'), 'w') as f:
|
|
json.dump(metadata, f)
|
|
|
|
base_image.save(out_path)
|
|
saved_base = True
|
|
|
|
except ValueError as e:
|
|
print(f'Error: {e}')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|