From fef12f235d63e1ec5e71380e6b19b01bd4bb87f0 Mon Sep 17 00:00:00 2001 From: Tetiana Yemelianenko <tyemel.mzeom@gmail.com> Date: Tue, 17 Sep 2024 13:53:35 +0000 Subject: [PATCH] Replace create_owl_dataset.py --- create_owl_dataset.py | 223 +++++++++++++++++++++++++----------------- 1 file changed, 134 insertions(+), 89 deletions(-) diff --git a/create_owl_dataset.py b/create_owl_dataset.py index 28caf1b..e0bf99e 100644 --- a/create_owl_dataset.py +++ b/create_owl_dataset.py @@ -9,20 +9,19 @@ import torch import numpy as np from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD -TOP_COUNT = 30 +TOP_COUNT = 10 feature_dim = 512 -parent_dir = 'path_to_the_main_dir' +parent_dir = '/home/tetiana/owl/' #path to the directory with images annotated on image level -parentpath = os.path.join(parent_dir, 'dataset/') -#print(parentpath) +parentpath = os.path.join(parent_dir, 'test/') -#directory in which we save selected images -owlpath = os.path.join(parent_dir, 'owl_dataset') +#directory in wich we save selected images +owlpath = os.path.join(parent_dir, 't_dataset') ext = '.jpg' #"Image file extension [.jpg or .png]" #path to the directory with non-annotated data -base_dir = 'path_to_the_non_annotated_dataset' +base_dir = '/home/tetiana/ty/projects/Archive_/' if not os.path.isdir(owlpath): os.mkdir(owlpath) @@ -39,11 +38,11 @@ processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble") model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device) #path to the csv file with saved TOP 50 regions of interests with their coordinates pre-calculated for the each painting from WikiArt dataset -df = pd.read_csv("path_to_the_objectnesses_fle/objectness_wikiart.csv") +df = pd.read_csv("/home/tetiana/ty/projects/Archive_/objectness_wikiart_final_all_new.csv") #path to the ANNOY index file for WikiArt dataset with information about similarity of the objects in the images t = AnnoyIndex(feature_dim, metric='angular') -t.load('path_to_the_annoy_index/annoy_wikiart.ann') +t.load('/home/tetiana/ty/projects/Archive_/annoy_new.ann') #receive similar images using ANNOY @@ -152,10 +151,10 @@ def center_to_corners_format(box): bboxes_corners = [x1, y1, x2, y2] return bboxes_corners -def rescale_owl(raw_image, box): +def rescale_owl(box, width, height): # rescale coordinates - img_h = raw_image.height - img_w = raw_image.width + img_h = height + img_w = width width_ratio = 1 height_ratio = 1 @@ -180,6 +179,103 @@ def create_new_name(counter): return fname +def to_yolo_format(box, width, height): + """ + Parameters + ---------- + box: bounding box to convert + wifth, height of the image + + ---------- + function convert from OWL to YOLO format + """ + box = center_to_corners_format(box) + box = rescale_owl(box, width, height) + x = (box[0] + box[2]) / 2 / width + y = (box[1] + box[3]) / 2 / height + w = (box[2] - box[0]) / width + h = (box[3] - box[1]) / height + return x, y, w, h + +def find_for_label(label, folder, imgpaths, threshold=0.4): + """ + Parameters + ---------- + label: current label for annotation + folder: current folder with images annotated on image level + imgpaths: list with paths to the images + threshold: current threshold + ---------- + function for creating annotations using files from folder annotated on image level + """ + annot = [] + count = 0 + iter = 0 + + for file in imgpaths: + iter += 1 + print(iter) + try: + boxes, query_embeddings = detectobject(file, label, threshold) + except: + print(file) + continue + #if current type of object wasn't detected with OWL, we skip this image + if len(boxes) == 0: + continue + + #for all found objects we use their embeddings to find similar objects in WikiArt dataset using ANNOY + for j in range(len(query_embeddings)): + #search objects similar to found embeddings using ANNOY + similar_img_ids, distances = get_similar_images_annoy(query_embeddings[j]) + df_selected = df.iloc[similar_img_ids] + #path to the image + similar_images = list(df_selected['file_path']) + #coordinates of the bounding box for the object + cxs = list(df_selected['cx']) + cys = list(df_selected['cy']) + ws = list(df_selected['w']) + hs = list(df_selected['h']) + + for k in range(TOP_COUNT): + found_file = os.path.join(base_dir, similar_images[k]) + raw_image = Image.open(found_file) + width, height = raw_image.size + raw_image.close() + + try: + found_boxes, _ = detectobject(found_file, label, threshold) + except: + print(found_file) + + if len(found_boxes) == 0: + continue + + box = center_to_corners_format([cxs[k], cys[k], ws[k], hs[k]]) + + max_iou = 0 + for p in range(len(found_boxes)): + iou = calculate_iou(box, center_to_corners_format(found_boxes[p])) + if iou > max_iou: + max_iou = iou + + if max_iou > 0.8: + #save found boundary boxes in YOLO format + x,y,w,h = to_yolo_format([cxs[k], cys[k], ws[k], hs[k]], width, height) + annot.append([found_file, folder, x, y, w, h]) + count += 1 + + #save boxes for the request image too + raw_image = Image.open(file) + width, height = raw_image.size + raw_image.close() + x,y,w,h = to_yolo_format(boxes[j], width, height) + annot.append([file, folder, x, y, w, h]) + count += 1 + + return annot, count + + #creation of annotations def owl_annoy_annotation(labels, folders): """ @@ -199,78 +295,25 @@ def owl_annoy_annotation(labels, folders): imgpaths = glob(parentpath + folders[i] + "/*" + ext) label = labels[i] folder = folders[i] - print(label) - iter = 0 - #for all images of the current label - for file in imgpaths: - iter += 1 - print(iter) - try: - boxes, query_embeddings = detectobject(file, label, 0.4) - except: - print(file) - continue - #if current type of object wasn't detected with OWL, we skip this image - if len(boxes) == 0: - continue - - #for all found objects we use their embeddings to find similar objects in WikiArt dataset using ANNOY - for j in range(len(query_embeddings)): - #search objects similar to found embeddings using ANNOY - similar_img_ids, distances = get_similar_images_annoy(query_embeddings[j]) - df_selected = df.iloc[similar_img_ids] - #path to the image - similar_images = list(df_selected['file_path']) - #coordinates of the bounding box for the object - cxs = list(df_selected['cx']) - cys = list(df_selected['cy']) - ws = list(df_selected['w']) - hs = list(df_selected['h']) - - for k in range(TOP_COUNT): - found_file = os.path.join(base_dir, similar_images[k]) - - raw_image = Image.open(found_file) - width, height = raw_image.size - raw_image.close() - - try: - found_boxes, _ = detectobject(found_file, label, 0.4) - except: - print(found_file) - - if len(found_boxes) == 0: - continue - - box = center_to_corners_format([cxs[k], cys[k], ws[k], hs[k]]) - - max_iou = 0 - for p in range(len(found_boxes)): - iou = calculate_iou(box, center_to_corners_format(found_boxes[p])) - if iou > max_iou: - max_iou = iou - - if max_iou > 0.8: - #save found boundary boxes in YOLO format - box = center_to_corners_format([cxs[k], cys[k], ws[k], hs[k]]) - box = rescale_owl(raw_image, box) - x = (box[0] + box[2]) / 2 / width - y = (box[1] + box[3]) / 2 / height - w = (box[2] - box[0]) / width - h = (box[3] - box[1]) / raw_image.height - annotations.append([found_file, folder, x, y, w, h]) - raw_image.close() - #save boxes for the request image too - raw_image = Image.open(file) - width, height = raw_image.size - box = center_to_corners_format(boxes[j]) - box = rescale_owl(raw_image, box) - x = (box[0] + box[2]) / 2 / width - y = (box[1] + box[3]) / 2 / height - w = (box[2] - box[0]) / width - h = (box[3] - box[1]) / raw_image.height - raw_image.close() - annotations.append([file, folder, x, y, w, h]) + print('Current label: ' + label) +# request_count = round(len(imgpaths)*1.2, 0) + request_count = 45 + selected_threshold = 0.4 + + not_enough_count = True + while not_enough_count: + print('next round') + cur_annot, found_count = find_for_label(label, folder, imgpaths, selected_threshold) + print('found on this round: ' + str(found_count)) + + if found_count < request_count: + selected_threshold -= 0.1 + if selected_threshold < 0.2: + not_enough_count = False + else: + not_enough_count = False + + annotations += cur_annot #get rid of from the duplicates and sort by file name the final list annotations = [list(x) for x in set(tuple(row) for row in annotations)] @@ -309,11 +352,13 @@ def owl_annoy_annotation(labels, folders): #list of labels -labels = ['an apple', 'a banana', 'a butterfly', 'a boat', 'a cat', 'a cow', 'a crucifixion', 'a deer', 'a dog', 'a white dove', 'an eagle', 'a horse', 'a monkey', > - 'a nude', 'a rooster', 'a serpent', 'a skull', 'a sheep', 'a swan', 'a trumpet'] - +#labels = ['an apple', 'a banana', 'a butterfly', 'a boat', 'a cat', 'a cow', 'a crucifixion', 'a deer', 'a dog', 'a white dove', 'an eagle', 'a horse', 'a monkey', > +# 'a nude', 'a rooster', 'a serpent', 'a skull', 'a sheep', 'a swan', 'a trumpet'] +labels = ['a dog', 'a swan'] #name of folders with previously collected request images annotated on image level -folders = ['apple', 'banana', 'butterfly', 'boat', 'cat', 'cow', 'crucifixion', 'deer', 'dog', 'dove', 'eagle', 'horse', 'monkey', 'orange', - 'nude', 'rooster', 'serpent', 'skull', 'sheep', 'swan', 'trumpet'] +#folders = ['apple', 'banana', 'butterfly', 'boat', 'cat', 'cow', 'crucifixion', 'deer', 'dog', 'dove', 'eagle', 'horse', 'monkey', 'orange', +# 'nude', 'rooster', 'serpent', 'skull', 'sheep', 'swan', 'trumpet'] + +folders = ['dog', 'swan'] owl_annoy_annotation(labels, folders) -- GitLab