app.py

# -*- coding: utf-8 -*-
# Copyright (c) Louis Brulé Naudet. All Rights Reserved.
# This software may be used and distributed according to the terms of License Agreement.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gradio as gr

try:
    import detectron2
    
except:
    import os 
    os.system('pip install git+https://github.com/louisbrulenaudet/detectron2')

import cv2
import torch

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog

model = {
    "name": "Sealife detection",
    "model_path": "./model_final.pth",
    "classes": ["creatures", "fish", "jellyfish", "penguin", "puffin", "shark", "starfish", "stingray"],
    "cfg": None,
    "metadata": None
}

model["cfg"] = get_cfg()
model["cfg"].merge_from_file("./configs/faster_rcnn_R_50_FPN_3x.yaml")
model["cfg"].MODEL.ROI_HEADS.NUM_CLASSES = len(model["classes"])
model["cfg"].MODEL.WEIGHTS = model["model_path"]

model["metadata"] = MetadataCatalog.get(model["name"])
model["metadata"].thing_classes = model["classes"]

if not torch.cuda.is_available():
    model["cfg"].MODEL.DEVICE = "cpu"


def inference(image, threshold):
    """
    Perform inference using a provided image and a pre-trained model, generating visual predictions.

    Parameters
    ----------
    image : numpy.ndarray
        The input image for which inference needs to be performed. It should be in BGR format.

    min_score : float
        The minimum confidence score threshold for predictions.

    model : detectron2.engine.DefaultPredictor
        The pre-trained model used for inference.

    Returns
    -------
    numpy.ndarray
        An image with drawn instance predictions generated by the model.
    """
    global model

    # Model expects BGR
    im = image[:, :, ::-1]

    # Set score threshold
    model["cfg"].MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold

    # Use the provided model
    predictor = DefaultPredictor(model["cfg"])
    outputs = predictor(im)

    v = Visualizer(im, model["metadata"], scale=1.2)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

    return out.get_image()[:, :, ::-1]


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", label="Input Image")
            threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, label="Threshold value")
            inference_button = gr.Button("Submit")
  
        image_output = gr.Image(type="pil", label="Output")

    inference_button.click(fn=inference, inputs=[image_input, threshold], outputs=image_output)

demo.launch()