Spaces:
Sleeping
Sleeping
| 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 | |
| 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() | |