import cv2
import numpy as np
from flask import Flask, jsonify, request
import os

app = Flask(__name__)

@app.route("/")
def indexfun():
  return "welcome"

@app.route("/predict", methods=['POST'])
def predictfun():
  net = cv2.dnn.readNet("yolov4-custom_2000.weights", "yolov4-custom.cfg")

  file = request.files['file']
  image_path = "./images/" + file.filename
  file.save(image_path)

  value_x = 0
  value_y = 0
  value_w = 0
  value_h = 0

  img = cv2.imread(image_path)
  height, width, _ = img.shape

  blob = cv2.dnn.blobFromImage(img, 1/255, (416,416), (0,0,0), swapRB=True, crop=False)
  net.setInput(blob)

  output_layers_names = net.getUnconnectedOutLayersNames()
  layerOutputs = net.forward(output_layers_names)

  boxes = []
  class_ids = []
  confidences = []

  for output in layerOutputs:
    for detection in output:
      scores = detection[5:]
      class_id = np.argmax(scores)
      confidence = scores[class_id]

      if confidence > 0.5:
        center_x = int(detection[0]*width)
        center_y = int(detection[1]*height)
        w = int(detection[2]*width)
        h = int(detection[3]*height)

        x = int(center_x - w/2)
        y = int(center_y - h/2)

        boxes.append([x, y, w, h])
        confidences.append((float(confidence)))
        class_ids.append(class_id)

  indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)

  if(len(boxes)>0):

    for i in indexes.flatten():
      x, y, w, h = boxes[i]
      #[70, 69, 116, 194]

      value_x = x
      value_y = y
      value_w = x+w
      value_h = y+h

    os.remove(image_path)

    return jsonify({
      'detected' : 1,
      'x': value_x,
      'w': value_w,
      'y': value_y,
      'h': value_h,
    })

  else:
    os.remove(image_path)
    return jsonify({
      'detected' : 0,
      'x': 0,
      'w': 0,
      'y': 0,
      'h': 0,
    })

if __name__ == "__main__":
  app.run(debug=True)
