app.py

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__), "yolo-lightnet")
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()