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
예제 #2
0
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]