Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import subprocess | |
| import json | |
| import os | |
| from PIL import Image | |
| darknetpath = os.path.join(os.path.dirname(__file__), "darknet") | |
| darknet_executable = os.path.join(darknetpath, "darknet") | |
| modelpath = os.path.join(os.path.dirname(__file__), "models") | |
| models = {} | |
| modellist = [] | |
| def init(): | |
| subprocess.run( | |
| "make", | |
| cwd=darknetpath, | |
| ) | |
| # subprocess.run( | |
| # "git lfs install", | |
| # cwd=modelpath, | |
| # ) | |
| # subprocess.run( | |
| # "git lfs pull", | |
| # cwd=modelpath, | |
| # ) | |
| global models | |
| models = json.load(open(os.path.join(modelpath, "path.json"))) | |
| global modellist | |
| modellist = list(models.keys()) | |
| def darknet_command(model, img, thresh=0.25): | |
| return [ | |
| darknet_executable, | |
| "detector", | |
| "test", | |
| model["data"], | |
| model["cfg"], | |
| model["weights"], | |
| img, | |
| "-thresh", | |
| str(thresh), | |
| ] | |
| def predict(model, img): | |
| input_path = os.path.join(modelpath, "input.jpg") | |
| output_path = os.path.join(modelpath, "predictions.jpg") | |
| img.save(input_path) | |
| model = models[model]["640x640"] | |
| command = darknet_command(model, input_path) | |
| subprocess.run(command, cwd=modelpath) | |
| return Image.open(output_path) | |
| if __name__ == "__main__": | |
| init() | |
| iface = gr.Interface( | |
| predict, | |
| inputs=[ | |
| gr.Dropdown(modellist, label="Model"), | |
| gr.Image(type="pil", label="Input Image"), | |
| ], | |
| outputs=gr.Image(type="pil", label="Output Image"), | |
| title="Yolo-lightnet", | |
| description="Yolo-lightnet is a lightweight version of Yolo. It removes the heavy layers of Yolo and replaces them with lightweight layers. This makes it faster and more efficient.", | |
| # examples=[ | |
| # [ | |
| # "driving", | |
| # "car.jpg", | |
| # ], | |
| # [ | |
| # "head_body", | |
| # "human.jpg", | |
| # ], | |
| # ], | |
| ) | |
| iface.launch() | |