import glob
import os
import shutil

import gradio as gr
import yolov5
from PIL import Image as im
model = yolov5.load('model/dango.pt')

# てすと


def inference(gr_input):
   # set model parameters
   model.conf = 0.45  # NMS confidence threshold
   model.iou = 0.45  # NMS IoU threshold
   model.agnostic = False  # NMS class-agnostic
   model.multi_label = False  # NMS multiple labels per box
   model.max_det = 1  # maximum number of detections per image
   results = model(gr_input)
   output_folder = "results"
   shutil.rmtree(output_folder)
   if not os.path.exists(output_folder):
      os.makedirs(output_folder)
   # ディレクトリを空にする

   # parse results
   predictions = results.pred[0]
   boxes = predictions[:, :4]  # x1, y1, x2, y2
   scores = predictions[:, 4]
   categories = predictions[:, 5]

   # boxesからx1,y1,x2,y2を取り出す
   x1 = boxes[:, 0].tolist()
   y1 = boxes[:, 1].tolist()
   x2 = boxes[:, 2].tolist()
   y2 = boxes[:, 3].tolist()

   # 4つの座標のうち、一つでも入っていなかったら、その画像はスキップ
   if x1 == [] or y1 == [] or x2 == [] or y2 == []:
      return

   x1 = int(x1[0])
   y1 = int(y1[0])
   x2 = int(x2[0])
   y2 = int(y2[0])
   img = im.open(gr_input)
   img_crop = img.crop((x1, y1, x2, y2))
   img_name = os.path.basename(gr_input)
   # pngで保存

   results.save(save_dir='results/')
   img_crop.save(f"{output_folder}/crop_{img_name}", quality=95)
   img_list = glob.glob("results/*.jpg")
   img_list.append(f"{output_folder}/crop_{img_name}")
   return img_list


iface = gr.Interface(
    fn=inference,
    inputs=gr.File(label="画像を選択"),
    outputs=gr.Gallery(label="結果"),
    title="YOLO"
)

if __name__ == "__main__":
   iface.launch()