import gradio as gr

import os
import sys
import numpy as np
import skimage.io
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from pycocotools import mask as maskUtils
import coco
from mrcnn.evaluate import build_coco_results, evaluate_coco
from mrcnn.dataset import MappingChallengeDataset
from mrcnn import visualize
from mrcnn.config import Config
from mrcnn import model as modellib, utils
import warnings
warnings.filterwarnings("ignore")



PRETRAINED_MODEL_PATH = os.path.join("data/","pretrained_weights.h5")
RETRAINED_MODEL_PATH = os.path.join("data/", "mask_rcnn_final.h5")
MODEL_DIR = os.path.join("logs")



class InferenceConfig(coco.CocoConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    NUM_CLASSES = 1 + 1  # 1 Background + 1 Building
    IMAGE_MAX_DIM=320
    IMAGE_MIN_DIM=320
    NAME = "crowdai-mapping-challenge"
config = InferenceConfig()



model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
model_path = PRETRAINED_MODEL_PATH
model.load_weights(model_path, by_name=True)
model.keras_model._make_predict_function()

class_names = ['BG', 'building'] # In our case, we have 1 class for the background, and 1 class for building

def binary_mask_area(mask):
    pixels = np.count_nonzero(np.float32(mask))
    one_pixel_area = 0.25 #0.25 meter^2
    return pixels*0.25

def classify_image(img):
    #random_image = skimage.io.imread(img)
    
    predictions = model.detect([img]*config.BATCH_SIZE, verbose=1) # We are replicating the same image to fill up the batch_size
    p = predictions[0]
    image = visualize.display_instances(img, p['rois'], p['masks'], p['class_ids'], class_names, p['scores'])

    mask = p['masks']
    mask = mask.astype(np.uint32)
    #mask.shape[2]
    final_mask = np.zeros((320, 320, 1), dtype = np.uint32)
    for c in range(mask.shape[2]):
        final_mask = np.bitwise_or(final_mask, mask[:, :, c].reshape(320, 320, 1))
        #final_mask.shape
    
    pixel_area = binary_mask_area(final_mask)
    msg = str(pixel_area) + " square meters"
    return (np.array(image), msg)

image = gr.inputs.Image(shape=(320, 320), type='numpy')
examples = ['test0.jpg', 'img2.png', 'img3.png', 'img4.png']

intf = gr.Interface(fn=classify_image, inputs=gr.Image(type='numpy'), outputs=['image', 'text'], examples=examples)
intf.launch()