Create app.py

app.py

@@ -0,0 +1,105 @@
+import json
+import random
+import spaces
+
+import gradio as gr
+import numpy as np
+import onnxruntime
+import torch
+from PIL import Image, ImageColor
+from torchvision.utils import draw_bounding_boxes
+import rfdetr.datasets.transforms as T
+
+
+def process_categories() -> tuple:
+    with open("categories.json") as fp:
+        categories = json.load(fp)
+
+    category_id_to_name = {d["id"]: d["name"] for d in categories}
+
+    random.seed(42)
+    color_names = list(ImageColor.colormap.keys())
+    sampled_colors = random.sample(color_names, len(categories))
+    rgb_colors = [ImageColor.getrgb(color_name) for color_name in sampled_colors]
+    category_id_to_color = {category["id"]: color for category, color in zip(categories, rgb_colors)}
+
+    return category_id_to_name, category_id_to_color
+
+
+
+def draw_predictions(boxes, labels, scores, img, score_threshold=0.5):
+    imgs_list = []
+    label_id_to_name, label_id_to_color = process_categories()
+
+    mask = scores > score_threshold
+    boxes_filtered = boxes[mask]
+    labels_filtered = labels[mask]
+    scores_filtered = scores[mask]
+
+    label_names = [label_id_to_name[int(i)] for i in labels_filtered]
+    colors = [label_id_to_color[int(i)] for i in labels_filtered]
+
+    img_bbox = draw_bounding_boxes(
+        img,
+        boxes=torch.from_numpy(boxes_filtered),
+        labels=[f"{name}: {score:.2f}" for name, score in zip(label_names, scores_filtered)],
+        colors=colors,
+        width=4
+    )
+    imgs_list.append(img_bbox.permute(1, 2, 0).numpy())  # convert to HWC for Gradio
+
+    return imgs_list
+
+
+@spaces.CPU(duration=20)
+def inference(image_path, model_name, bbox_threshold):
+    transforms = T.Compose([
+    T.SquareResize([1120]),
+    T.ToTensor(),
+    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+    image = Image.open(image_path).convert("RGB")
+    tensor_img, _ = transforms(image, None)
+    tensor_img = tensor_img.unsqueeze(0)
+
+    ort_inputs = {
+        'input': tensor_img.cpu().numpy()
+    }
+    
+    model_path = "/home/datsplit/FashionVeil/models/rfdetr/onnx-models/rfdetrl_finetuned_fashionveil.onnx"
+
+    
+    sess_options = onnxruntime.SessionOptions()
+    sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
+    ort_session = onnxruntime.InferenceSession(
+        model_path,
+        providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+        sess_options=sess_options
+    )
+
+    ort_inputs = {ort_session.get_inputs()[0].name: img_transformed}
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    boxes, labels, scores = ort_outs
+    return draw_predictions(boxes, labels, scores, torch.from_numpy(np.array(img)), score_threshold=bbox_threshold)
+
+
+
+title = "FashionUnveil - Demo"
+description = r"""This is the demo of the research project <a href="https://github.com/DatSplit/FashionVeil">FashionUnveil</a>. Upload your image for inference."""
+
+demo = gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.Image(type="filepath", label="Input Image"),
+        gr.Dropdown(["RF-DETR-L"], value="RF-DETR-L", label="Model"),
+        gr.Slider(value=0.5, minimum=0.0, maximum=0.9, step=0.05, label="BBox threshold"),
+    ],
+    outputs=gr.Gallery(label="Output", preview=True, height=500),
+    title=title,
+    description=description,
+)
+
+if __name__ == "__main__":
+    demo.launch()