app.py

#!/usr/bin/env python
from tqdm.notebook import tqdm as notebook_tqdm
from detectron2.config import LazyConfig, instantiate
from detectron2.checkpoint import DetectionCheckpointer
from demo.demo import VisualizationDemo, get_parser
import numpy as np
from detectron2.data.detection_utils import read_image
import matplotlib.pyplot as plt
import cv2
import os
import gradio as gr
from PIL import Image

config_file = 'demo-confs/odor_demo.py'
ckpt_pth = 'focaldino_ep18.pth'

try:
    cfg = LazyConfig.load(config_file)
except AssertionError as e:
    if str(e).startswith('Dataset '):
        pass
    else:
        raise e
print("Loading the model...")
model = instantiate(cfg.model)
model.to("cpu")
checkpointer = DetectionCheckpointer(model)
checkpointer.load(ckpt_pth)
model.eval();
print("Model loaded. Let's predict.")


def treat_grayscale(img):
    if len(img.shape) == 2:
        return np.stack((img,)*3, axis=-1)
    else:
        return img

def show_preds(input_image):

    demo = VisualizationDemo(
        model=model,
        min_size_test=800,
        max_size_test=1333,
        img_format='RGB',
        metadata_dataset='odor_test'
    )
    img = treat_grayscale(input_image)
    predictions, visualized_output = demo.run_on_image(img, 0.25)
    visualized_rgb_image = visualized_output.get_image()
    visualized_rgb_image = visualized_rgb_image[:,:,::-1]
    return visualized_rgb_image


gr_interface = gr.Interface(
    fn=show_preds,
    inputs=["image"],
    #outputs=[gr.outputs.Image(type="pil")],
    outputs=[gr.Image(type="pil")],
    title="ODOR Object Detector",
    description="A DINO model with focalnet architecture that detects objects relating to smells. Upload an image or click an example image below to use.",
)
gr_interface.launch()