def detect(checkpoint, input): score_prediction = 0.5 chip_size = 300 #Histogram Equalize the image img = cv2.imread(input, 0) histImg = cv2.equalizeHist(img) clahe3 = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) cl13 = clahe3.apply(histImg) #Parse and chip images arr = np.array(np.expand_dims(cl13, axis=2)) chip_size = (chip_size,chip_size) images = chip_image(arr,chip_size) #generate detections boxes, scores, classes = generate_detections(checkpoint,images) #Process boxes to be full-sized width,height,_ = arr.shape cwn,chn = (chip_size) wn,hn = (int(width/cwn),int(height/chn)) num_preds = 250 bfull = boxes[:wn*hn].reshape((wn,hn,num_preds,4)) b2 = np.zeros(bfull.shape) b2[:,:,:,0] = bfull[:,:,:,1] b2[:,:,:,1] = bfull[:,:,:,0] b2[:,:,:,2] = bfull[:,:,:,3] b2[:,:,:,3] = bfull[:,:,:,2] bfull = b2 bfull[:,:,:,0] *= cwn bfull[:,:,:,2] *= cwn bfull[:,:,:,1] *= chn bfull[:,:,:,3] *= chn for i in range(wn): for j in range(hn): bfull[i,j,:,0] += j*cwn bfull[i,j,:,2] += j*cwn bfull[i,j,:,1] += i*chn bfull[i,j,:,3] += i*chn bfull = bfull.reshape((hn*wn,num_preds,4)) #get the file names img_name = os.path.basename(input).split('.')[0] model_name = os.path.basename(checkpoint).split('_')[1].split('.')[0] file_name = img_name + '_' + model_name + '_preds.txt' with open(file_name,'w') as f: for i in range(bfull.shape[0]): for j in range(bfull[i].shape[0]): #box should be xmin ymin xmax ymax box = bfull[i,j] class_prediction = classes[i,j] score_prediction = scores[i,j] f.write('%d %d %d %d %d %f \n' % \ (box[0],box[1],box[2],box[3],int(class_prediction),score_prediction)) # Set up the dataframe from the txt file data = pandas.read_csv(file_name, header=-1) data.columns = ['init'] new = data["init"].str.split(" ", expand=True) data["xmin"]= new[0].astype(int) data["ymin"]= new[1].astype(int) data["xmax"]= new[2].astype(int) data["ymax"]= new[3].astype(int) data["class_prediction"]= new[4].astype(int) data["score_prediction"]= new[5].astype(float) data.drop(["init"], axis=1, inplace = True) # Confidence data_con = data[data.score_prediction >= score_prediction] # Set up the picture and draw the boxes orgImg = cv2.imread(input) # Save the image correctly cv2.imwrite(img_name + '_' + model_name + '.png', orgImg) name = img_name + '_' + model_name + '_' + 'preds.txt' imagename = img_name + '_' + model_name + '.png' return name, imagename
if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("-c","--checkpoint", default='pbs/model.pb', help="Path to saved model") parser.add_argument("-cs", "--chip_size", default=300, type=int, help="Size in pixels to chip input image") parser.add_argument("input", help="Path to test chip") parser.add_argument("-o","--output",default="predictions.txt",help="Filepath of desired output") args = parser.parse_args() #Parse and chip images arr = np.array(Image.open(args.input)) chip_size = (args.chip_size,args.chip_size) images = chip_image(arr,chip_size) print(images.shape) #generate detections boxes, scores, classes = generate_detections(args.checkpoint,images) #Process boxes to be full-sized width,height,_ = arr.shape cwn,chn = (chip_size) wn,hn = (int(width/cwn),int(height/chn)) num_preds = 250 bfull = boxes[:wn*hn].reshape((wn,hn,num_preds,4)) b2 = np.zeros(bfull.shape) b2[:,:,:,0] = bfull[:,:,:,1] b2[:,:,:,1] = bfull[:,:,:,0] b2[:,:,:,2] = bfull[:,:,:,3] b2[:,:,:,3] = bfull[:,:,:,2] bfull = b2
empty_image_idx = [] for idx, image in enumerate(im): # skip chips that do not have buildings, avoid feeding them into inference if len(box_chip[idx]) == 0 or (len(box_chip[idx]) == 1 and np.all(box_chip[idx] == 0)): empty_image_idx.append(idx) k = k + 1 else: images_list.append(image) images = np.array(images_list) images.astype(np.uint8) print('number of images without bboxes: ', k) print('images shape: ', images.shape) i = 0 boxes_pred, scores_pred, classes_pred = generate_detections( args.checkpoint, images) # debug for idx, image in enumerate(im): if idx in set(empty_image_idx): box = np.zeros((1, num_preds, 4)) score = np.zeros((1, num_preds)) clss = np.zeros((1, num_preds)) boxes.append(box) scores.append(score) classes.append(clss) else: #continue box_pred, score_pred, cls_pred = boxes_pred[i], scores_pred[ i], classes_pred[i]