import gradio as gr
import os
import torch

os.system('git clone https://github.com/WongKinYiu/yolov7')


def detect(inp):
  os.system('python ./yolov7/detect.py --weights best.pt --conf 0.25 --img-size 640 --source f{inp} --project ./yolov7/runs/detect ')
  otp=inp.split('/')[2]
  return f"./yolov7/runs/detect/exp/*"
   
  #f"./yolov7/runs/detect/exp/{otp}"
  
  
  


def custom(path_or_model='path/to/model.pt', autoshape=True):
    """custom mode
    Arguments (3 options):
        path_or_model (str): 'path/to/model.pt'
        path_or_model (dict): torch.load('path/to/model.pt')
        path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
    Returns:
        pytorch model
    """
    model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model  # load checkpoint
    if isinstance(model, dict):
        model = model['ema' if model.get('ema') else 'model']  # load model
    
    hub_model = Model(model.yaml).to(next(model.parameters()).device)  # create
    hub_model.load_state_dict(model.float().state_dict())  # load state_dict
    hub_model.names = model.names  # class names
    if autoshape:
        hub_model = hub_model.autoshape()  # for file/URI/PIL/cv2/np inputs and NMS
    device = select_device('0' if torch.cuda.is_available() else 'cpu')  # default to GPU if available
    return hub_model.to(device)
    
model = custom(path_or_model='best.pt')

def detect1(inp):
 #g = (size / max(inp.size))  #gain
 #im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize 
 results = model(inp,size=640)  # inference
 results.render()  # updates results.imgs with boxes and labels
 return Image.fromarray(results.imgs[0])  

inp = gr.inputs.Image(type="filepath", label="Input")
#output=gr.outputs.Image(type="pil", label="Output Image")
output = gr.outputs.Image(type="filepath", label="Output")
#.outputs.Textbox()

io=gr.Interface(fn=detect1, inputs=inp, outputs=output, title='Pot Hole Detection With Custom YOLOv7 ',
#examples=[["Examples/img-300_jpg.rf.6b7b035dff1cda092ce3dc22be8d0135.jpg"]]
)
#,examples=["Examples/img-300_jpg.rf.6b7b035dff1cda092ce3dc22be8d0135.jpg"]
io.launch(debug=True,share=False)