Spaces:
Running
Running
File size: 7,121 Bytes
1d64201 d1af37b 1d64201 bce365b 1d64201 b09ff42 1d64201 0872418 bce365b 1d64201 b09ff42 1d64201 f8b59fe 1d64201 0872418 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 |
import os
import cv2
from pycocotools.coco import COCO
from detectron2.engine import DefaultPredictor
import matplotlib.pyplot as plt
import torch
from datetime import datetime
from rcnn_config import write_config
from from_root import from_root
from rcnn_model.preprocessing.cleaning_single_image import preprocess_image
from rcnn_model.utils.floorplan_vectorizer_utils import draw_from_coco, bitmask_to_polygon
from rcnn_model.extraction.annotation_builder import AnnotationBuilder as AnnBuild
results_directory = "/tmp/results/"
sample_data_directory = "rcnn_model/sample/"
model_directory = "/tmp/output/"
def main(cfg,img_source_path, coco_dest_filename, val_img_dest_filename):
os.chdir(str(from_root()))
print("Inside main: " + img_source_path)
#configure model
cfg.DATALOADER.NUM_WORKERS = 1
cfg.SOLVER.IMS_PER_BATCH = 1
cfg.MODEL.WEIGHTS = os.path.join(model_directory, "model_final.pth")
cfg.MODEL.DEVICE = "cpu"
predictor = DefaultPredictor(cfg)
#run
prediction_runner(img_source_path, results_directory+coco_dest_filename, results_directory+val_img_dest_filename, predictor)
print("SAVED to "+results_directory+coco_dest_filename+" and "+results_directory+val_img_dest_filename)
### Main Runner ###
def prediction_runner(filename, coco_dest_path, val_img_dest_path, predictor, segmented_prediction = False, scale_factor = 0):
#set up annotation builder
ann_builder = instantiate_ann_build()
print("Filename inside prediction_runner: ", filename)
#set up image
initImg = preprocess_image(filename)
init_width = initImg.shape[1]
init_height = initImg.shape[0]
img_id = ann_builder.add_image(str(from_root(filename)), init_width, init_height)
print("cleaned")
#resize image
scaled_width = 800
scaled_height = 800
img = cv2.resize(initImg, (scaled_width,scaled_height))
if(scale_factor > 0):
img = cv2.resize(initImg, (0,0), fx=scale_factor, fy=scale_factor)
scaled_width = img.shape[1]
scaled_height = img.shape[0]
print("resized")
#run prediction
if(segmented_prediction):
run_segmented_prediction(ann_builder, predictor, img, img_id, scaled_width, scaled_height, scale_factor)
else:
run_prediction(ann_builder, predictor, img, img_id, init_width/scaled_width, init_height/scaled_height)
#save the file
ann_builder.save_file(str(from_root(coco_dest_path)))
#visualize
coco = COCO(str(from_root(coco_dest_path)))
draw_from_coco(0, coco, val_img_dest_path)
def instantiate_ann_build():
ann_builder = AnnBuild()
ann_builder.set_info("Generated Annotations from User Submitted Data","End User","N/A",datetime.now())
ann_builder.add_license("TBD", "TBD")
return ann_builder
def prediction_outputs_to_annotations(annotations, outputs, img_id, base_ann_id, x_offset=0, y_offset=0, h_scale_factor=1, v_scale_factor=1):
ann_id = base_ann_id
for i in range(0,len(outputs["instances"].to(torch.device("cpu")).pred_masks)):
mask = outputs["instances"].to(torch.device("cpu")).pred_masks[i]
class_id = outputs["instances"].to(torch.device("cpu")).pred_classes[i].item()
score = outputs["instances"].to(torch.device("cpu")).scores[i].item()
annotations.append(bitmask_to_polygon(ann_id, img_id, class_id, score, mask, x_offset=x_offset, y_offset=y_offset,scale_factor_width = h_scale_factor,scale_factor_height = v_scale_factor))
ann_id += 1
return annotations, ann_id
### Standard Prediction ###
def run_prediction(ann_builder, predictor, img, img_id, h_scale_factor, v_scale_factor):
outputs = predictor(img)
print("predicted")
annotations = []
annotations, annId = prediction_outputs_to_annotations(annotations, outputs, img_id, 0, 0, 0, h_scale_factor, v_scale_factor)
ann_builder.annotations = annotations
print("annotations converted")
### Segmented Prediction ###
def run_segmented_prediction(ann_builder, predictor, img, img_id, width, height, scale_factor, segment_size = 800):
#initialize
count = 1
subimg_dest_path = "models/rcnn/westmoor_check/subimgs/subimg_"
annotations = []
ann_id = 0
#iterate through segments
for xi in range(0, int(width/segment_size)+1):
for yi in range(0, int(height/segment_size)+1):
#calculate subimg area
h_base, v_base, h_boundary, v_boundary = get_subimg_area(xi,yi,width,height,segment_size)
#save subimgs
if(h_boundary > h_base and v_boundary > v_base):
subimg = img[v_base:v_boundary,h_base:h_boundary,:]
save_subimg(subimg, subimg_dest_path, count, h_boundary-h_base, v_boundary-v_base)
#get annotations
outputs = predictor(subimg)
annotations, ann_id = prediction_outputs_to_annotations(annotations, outputs, img_id, ann_id, h_base*1/scale_factor, h_base*1/scale_factor, 1/scale_factor, 1/scale_factor)
count += 1
ann_builder.annotations = annotations
def get_subimg_area(xi,yi,img_width,img_height,segment_size):
h_base = xi*segment_size
v_base = yi*segment_size
h_boundary = (xi+1)*segment_size
if(h_boundary >= img_width):
h_boundary = img_width-1
v_boundary = (yi+1)*segment_size
if(v_boundary >= img_height):
v_boundary = img_height-1
return h_base, v_base, h_boundary, v_boundary
def save_subimg(subimg, subimg_dest_path, count, width, height):
plt.figure(figsize=(width, height),dpi=1)
plt.imshow(subimg)
plt.axis('off')
plt.savefig(str(from_root(subimg_dest_path+str(count)+".png")),bbox_inches='tight',pad_inches=0)
plt.clf()
def default_sample(cfg, run_index = -1):
if(run_index >= 0):
main(cfg,sample_data_directory+"westmoor_floor_2_floorplan.png","westmoor_result_run_"+str(run_index)+".json","westmoor_result_run_"+str(run_index)+".png")
main(cfg,sample_data_directory+"REV2 ENSURE Layout - Springs at the Waterfront - PR-2023-3661 - CN0012630 - 11.02.2023 (1)-16.png","rev2_result_run_"+str(run_index)+".json","rev2_result_run_"+str(run_index)+".png")
main(cfg,sample_data_directory+"F1_original.png","cubicasa_result_run_"+str(run_index)+".json","cubicasa_result_run_"+str(run_index)+".png")
else:
main(cfg,sample_data_directory+"westmoor_floor_2_floorplan.png","westmoor_result.json","westmoor_result.png")
main(cfg,sample_data_directory+"REV2 ENSURE Layout - Springs at the Waterfront - PR-2023-3661 - CN0012630 - 11.02.2023 (1)-16.png","rev2_result.json","rev2_result.png")
main(cfg,sample_data_directory+"F1_original.png","cubicasa_result.json","cubicasa_result.png")
if __name__ == "__main__":
print("Inside Main Function Call")
cfg = write_config()
run_index=-1
#Sample run for a single image
main(cfg,sample_data_directory+"F1_original.png","cubicasa_result_run_"+str(run_index)+".json","cubicasa_result_run_"+str(run_index)+".png")
|