<filename>utils/Predict_Weed.py
import base64
import numpy as np
import tensorflow as tf
import cv2
from tqdm import tqdm
import json
import pickle
import re



model_path = './models/RCNN_crop_weed_classification_model.h5'
svm_model_path = './models/svm_classifier.pkl'
def get_weed_model():
    global model
    global svm_model
    model = tf.keras.models.load_model(model_path)
    model = tf.keras.models.Model(model.inputs,model.layers[-5].output)
    with open(svm_model_path,'rb') as svm:
        svm_model = pickle.load(svm)

print(" * Loading Keras model for weed detector...")
get_weed_model()


def iou_calc(bb1, bb2):

    true_xmin, true_ymin, true_width, true_height  = bb1
    bb_xmin, bb_ymin,  bb_width, bb_height = bb2

    true_xmax = true_xmin + true_width
    true_ymax = true_ymin + true_height
    bb_xmax = bb_xmin + bb_width
    bb_ymax = bb_ymin + bb_height

    #calculating area
    true_area = true_width * true_height
    bb_area   = bb_width * bb_height

    #calculating itersection cordinates
    inter_xmin = max(true_xmin , bb_xmin)
    inter_ymin = max(true_ymin , bb_ymin)
    inter_xmax = min(true_xmax , bb_xmax)
    inter_ymax = min(true_ymax , bb_ymax)

    if inter_xmax <= inter_xmin or inter_ymax <= inter_ymin:
        iou = 0


    else:
        inter_area = (inter_xmax - inter_xmin) * (inter_ymax - inter_ymin)


        iou = inter_area / (true_area + bb_area - inter_area)

    assert iou<=1
    assert iou>=0

    return iou

def predict_weed(img_path):
    confidence = 0.9
    iou_thresh = 0.1

    # appling selective search
    img = cv2.imread(img_path)
    cv2.setUseOptimized(True);
    ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()
    ss.setBaseImage(img)
    ss.switchToSelectiveSearchFast()
    rects = ss.process()
    sel_rects = rects[:2000]

    pred_crop = []
    pred_weed = []
    for index, rect in tqdm(enumerate(sel_rects)):

        x, y, w, h = rect
        roi = img[y:y + h, x:x + w, :]
        resized_roi = cv2.resize(roi, (224, 224)) / 255

        # Feature extraction

        feature = model.predict(resized_roi.reshape(-1, 224, 224, 3))

        # SVM prediction
        pred = svm_model.predict_proba(feature.reshape(-1, 4096))
        pred_lab = svm_model.predict(feature.reshape(-1, 4096))

        if pred_lab == 'crop' and np.max(pred) > confidence:
            pred_crop.append([list(rect), np.max(pred)])
        elif pred_lab == 'weed' and np.max(pred) > confidence:
            pred_weed.append([list(rect), np.max(pred)])

    final = []

    # Detection for crop class
    if len(pred_crop) != 0:
        pred_score_crop = [x[1] for x in pred_crop]
        pred_bb_crop = [x[0] for x in pred_crop]

        for i in range(len(pred_crop)):
            temp_bb, temp_score = pred_bb_crop.copy(), pred_score_crop.copy()
            if len(temp_bb) != 0:

                max_score_box = temp_bb[np.argmax(temp_score)]

                if [max_score_box, np.max(temp_score)] not in final:
                    final.append([max_score_box, np.max(temp_score), 'crop'])
                    index_should_del = []

                    for ind, other_bb in enumerate(temp_bb):
                        iou_score = iou_calc(max_score_box, other_bb)

                        # Non maximum suppression(nms)

                        if iou_score >= iou_thresh:
                            index_should_del.append(ind)

                    pred_bb_crop = []
                    pred_score_crop = []
                    for bb_index, bb_value in enumerate(temp_bb):
                        if bb_index not in index_should_del:
                            pred_bb_crop.append(bb_value)

                    for score_index, score_value in enumerate(temp_score):
                        if score_index not in index_should_del:
                            pred_score_crop.append(score_value)
                else:
                    continue

            else:
                break

    # Detection for weed class

    if len(pred_weed) != 0:
        pred_score_weed = [x[1] for x in pred_weed]
        pred_bb_weed = [x[0] for x in pred_weed]

        for i in range(len(pred_weed)):
            temp_bb, temp_score = pred_bb_weed.copy(), pred_score_weed.copy()
            if len(temp_bb) != 0:

                max_score_box = temp_bb[np.argmax(temp_score)]

                if [max_score_box, np.max(temp_score)] not in final:
                    final.append([max_score_box, np.max(temp_score), 'weed'])
                    index_should_del = []

                    for ind, other_bb in enumerate(temp_bb):
                        iou_score = iou_calc(max_score_box, other_bb)

                        if iou_score >= iou_thresh:
                            index_should_del.append(ind)

                    pred_bb_weed = []
                    pred_score_weed = []
                    for bb_index, bb_value in enumerate(temp_bb):
                        if bb_index not in index_should_del:
                            pred_bb_weed.append(bb_value)

                    for score_index, score_value in enumerate(temp_score):
                        if score_index not in index_should_del:
                            pred_score_weed.append(score_value)
                else:
                    continue

            else:
                break

    imOut = img.copy()
    crop_cls = 0
    weed_cls = 0
    for rect, score, cls in final:

        x, y, w, h = rect
        if cls == 'weed':
            color = (0, 0, 255)
            weed_cls = weed_cls + 1

        if cls == 'crop':
            color = (30, 255, 255)
            crop_cls = crop_cls + 1


        if crop_cls > weed_cls:
            percentage = round((crop_cls / (weed_cls + crop_cls)) * 100, 2)
            class_final = 'crop'
        if weed_cls > crop_cls:
            percentage = round((weed_cls / (weed_cls + crop_cls)) * 100, 2)
            class_final = 'weed'

        cv2.rectangle(imOut, (x, y), (x + w, y + h), color, 2)

        cv2.putText(imOut, cls + ':' + str(round(score * 100, 2)), (x, y - 8), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2,
                    cv2.LINE_AA)

    cv2.imwrite('prediction.jpeg', imOut)
    predicted = {'class': class_final, 'score': percentage}

    retval, buffer = cv2.imencode('.jpeg', imOut)

    base64_bytes = base64.b64encode(buffer)
    base64_str = base64_bytes.decode('utf-8')
    imaagedict = {'image': base64_str}

    decoded = base64.b64decode(base64_str)

    test_image = "testprediction.jpeg"
    with open(test_image, 'wb') as f:
        f.write(decoded)

    predicted.update(imaagedict)

    with open("weed-detection.json", "w") as outfile:
        json.dump(predicted, outfile)

    return predicted