Spaces:
Sleeping
Sleeping
File size: 9,901 Bytes
1d64201 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 |
import cv2
import numpy as np
import copy
from pycocotools.coco import COCO
import shapely
from shapely import geometry
import sys
import random
from datetime import datetime
from from_root import from_root
from annotation_builder import AnnotationBuilder as AnnBuild
import pylab
pylab.rcParams['figure.figsize'] = (128.0, 160.0)
from rcnn_model.utils.floorplan_vectorizer_utils import get_image_size, draw_from_coco
data_directory_root = str(from_root("dataset/"))+"/"
category_filter = [2]
image_sample_room_count_threshold = 4
min_sample_size = 400
max_sample_size = 800
samples_per_image = 30
def main():
sample_from_labelme2coco_dataset("train",data_directory_root+"annotations/","sample_data/","validation_images/")
sample_from_labelme2coco_dataset("val",data_directory_root+"annotations/","sample_data/","validation_images/")
### Core Sampling Logic ###
#sample from dataset cocofile created by labelme2coco
#dataset_name should only be "train", "val", or "dataset" based on labelme2coco's output naming conventions
def sample_from_labelme2coco_dataset(dataset_name,annotation_source_dir,sample_img_dest_dir,validation_img_dest_dir=""):
#initialize annbuilder
ann_builder = AnnBuild()
ann_builder.set_info("manual annotations of Inovonics and university provided data","inovonData","NA",datetime(2019,5,24))
ann_builder.add_license("TODO", "TODO")
coco = COCO(annotation_source_dir+dataset_name+".json")
print("Coco Loaded")
#reading
for img_id in coco.getImgIds():
take_samples_from_image(ann_builder, img_id, coco, sample_img_dest_dir)
#save
ann_builder.save_file(annotation_source_dir+dataset_name+"_sampled_data.json")
#validation images
if(validation_img_dest_dir != ""):
validation_coco = COCO(annotation_source_dir+dataset_name+"_sampled_data.json")
validation_images(dataset_name,validation_coco,validation_img_dest_dir)
def take_samples_from_image(ann_builder, img_id, coco, img_dest):
#set up image name
source_img_filename = coco.imgs[img_id]['file_name']
source_img = cv2.imread(data_directory_root+source_img_filename,cv2.IMREAD_COLOR)
img_name = source_img_filename[1:-4]
img_name = img_name[img_name.index("/"):]
img_name = img_name[1:]
#set up mirroring
mirrored_imgs = [source_img, np.fliplr(source_img), np.flipud(source_img), np.flipud(np.fliplr(source_img))]
mirror_tags = ["","_h","_v","_hv"]
print("Processing image "+str(img_id)+": "+img_name)
#run sampler
for m in range(0,len(mirrored_imgs)):
#load mirror of image
img = mirrored_imgs[m]
tag = mirror_tags[m]
mirrored_anns = mirror_coco_coordinates(coco,img_id,m)
if len(mirrored_anns) > 10:
#collect samples
for i in range(0,samples_per_image):
img_dest_path = data_directory_root+img_dest+img_name+tag+"_"+str(i)+".png"
take_sample(ann_builder, mirrored_anns, img, img_dest_path)
else:
img_dest_path = data_directory_root+img_dest+img_name+tag+".png"
take_full_image(ann_builder, mirrored_anns, img, img_dest_path)
def take_sample(ann_builder, annotations, img, img_dest_path):
#Take a random sample with at least a certain number of room bounding boxes overlapping
sample_annotations, cropped, room_count = random_sample_selection(annotations,img)
while(room_count < image_sample_room_count_threshold):
sample_annotations, cropped, room_count = random_sample_selection(annotations,img)
#sav the cropped image portion of the final sample
cv2.imwrite(img_dest_path, cropped)
print(" sample saved to "+img_dest_path)
sample_width, sample_height = get_image_size(img_dest_path)
sampled_img_id = ann_builder.add_image(img_dest_path, sample_width, sample_height)
#crop annotations
crop_area = define_crop_area(0,0,sample_width, sample_height)
add_cropped_annotations(ann_builder, sampled_img_id, sample_annotations, crop_area)
def take_full_image(ann_builder, annotations, img, img_dest_path):
#sav the cropped image portion of the final sample
cv2.imwrite(img_dest_path, img)
print(" whole image saved to "+img_dest_path)
width, height = get_image_size(img_dest_path)
img_id = ann_builder.add_image(img_dest_path, width, height)
for ann in annotations:
ann_builder.add_annotation(img_id, ann["category_id"], segmentation_to_polygon(ann["segmentation"]))
def random_sample_selection(annotations, img):
#get bounds of original image
init_width = len(img)
init_height = len(img[0])
#randomly select a rectangle
sample_x = random.randrange(0,init_width-min_sample_size-1)
sample_y = random.randrange(0,init_height-min_sample_size-1)
sample_width = random.randrange(min_sample_size,min(max_sample_size, init_width-sample_x))
sample_height = random.randrange(min_sample_size,min(max_sample_size, init_height-sample_y))
#create cropped image and offset annotation coordinates immediately for easier data transfer
cropped = img[sample_y:sample_y+sample_height,sample_x:sample_x+sample_width]
sampled_annotations, room_count = offset_annotation_coordinates(annotations,sample_x,sample_y,sample_width,sample_height)
#return values
return sampled_annotations, cropped, room_count
### Annotation Cropping ###
def define_crop_area(x, y, width, height):
cropped_area = geometry.Polygon([(x,y),
(x+width,y),
(x+width,y+height),
(x,y+height)])
return cropped_area
def add_cropped_annotations(ann_builder, img_id, annotations, cropped_area):
for ann in annotations:
#get intersecting area
poly = crop_polygon(ann["segmentation"],cropped_area,ann)
#handle convex rooms that weren't split
if(isinstance(poly,geometry.Polygon)):
ann_builder.add_annotation(img_id, ann["category_id"], poly)
#handle concave rooms that were split
elif(isinstance(poly,geometry.GeometryCollection) or isinstance(poly,geometry.MultiPolygon)):
for subpoly in poly.geoms:
if(isinstance(subpoly,geometry.Polygon)):
ann_builder.add_annotation(img_id, ann["category_id"], subpoly)
def crop_polygon(segmentation,crop_area,id):
#reformat into shapely geometry Polygon
poly = segmentation_to_polygon(segmentation)
#check shape validity (most common error is self overlapping)
if(not shapely.is_valid(poly)):
print(id)
print(shapely.is_valid_reason(poly))
return None
#check size and return
cropped_poly = shapely.intersection(poly,crop_area)
if(cropped_poly.area > 0):
return cropped_poly
else:
return None
def segmentation_to_polygon(segmentation):
points = np.array(segmentation[0])
points = points.reshape(int(len(segmentation[0])/2),2)
return geometry.Polygon([[p[0], p[1]] for p in points])
### Applying Geometry Modifications
def mirror_coco_coordinates(coco, img_id, mirroring_index):
#instantiate copy
original_annotations = coco.imgToAnns[img_id]
annotations = copy.deepcopy(original_annotations)
new_annotations = []
#get width
width = coco.imgs[img_id]['width']
height = coco.imgs[img_id]['height']
#apply mirroring
for ann in annotations:
if(ann["category_id"] in category_filter):
for i in range(0,len(ann['bbox'])):
apply_mirroring_to_coord(ann['bbox'], i, mirroring_index, width, height)
for i in range(0,len(ann['segmentation'][0])):
apply_mirroring_to_coord(ann['segmentation'][0], i, mirroring_index, width, height)
new_annotations.append(ann)
return new_annotations
def apply_mirroring_to_coord(array, index, mirroring_index, width, height):
if(index%2 == 0):
if(mirroring_index%2==1):
array[index] = width-array[index]
else:
if(mirroring_index>1):
array[index] = height-array[index]
def offset_annotation_coordinates(original_annotations,x_offset,y_offset,width,height):
#instantiate copy
annotations = copy.deepcopy(original_annotations)
new_annotations = []
room_count = 0
#apply offfset
for ann in annotations:
if(check_bounding_box_overlap(ann['bbox'], x_offset, y_offset, width, height)):
room_count += 1
for i in range(0,len(ann['bbox'])):
apply_offset_to_coord(ann['bbox'], i, x_offset, y_offset)
for i in range(0,len(ann['segmentation'][0])):
apply_offset_to_coord(ann['segmentation'][0], i, x_offset, y_offset)
new_annotations.append(ann)
return new_annotations, room_count
def apply_offset_to_coord(array, index, x_offset, y_offset):
if(index%2 == 0):
array[index] -= x_offset
else:
array[index] -= y_offset
def check_bounding_box_overlap(bbox, x_offset, y_offset, width, height):
boundary_threshold = 25 #so a row of rooms with a one pixel sliver within the area don't get counted
within_horizontal_bounds = max(bbox[0],bbox[2]) >= x_offset+boundary_threshold and min(bbox[0],bbox[2]) <= x_offset+width-boundary_threshold
within_vertical_bounds = max(bbox[1],bbox[3]) >= y_offset+boundary_threshold and min(bbox[1],bbox[3]) <= y_offset+height-boundary_threshold
return within_horizontal_bounds and within_vertical_bounds
### Validation Display ###
def validation_images(dataset_name,coco,validation_img_target_dir):
count = 1
for i in np.random.choice(coco.getImgIds(),8):
draw_from_coco(i,coco,data_directory_root+validation_img_target_dir+dataset_name+"_sampling_validation_"+str(count)+".png")
count+=1 |