Spaces:
Sleeping
Sleeping
| import argparse | |
| from functools import partial | |
| import cv2 | |
| import requests | |
| import os | |
| from io import BytesIO | |
| from PIL import Image | |
| import numpy as np | |
| from pathlib import Path | |
| import gradio as gr | |
| import warnings | |
| import torch | |
| import Equirec2Perspec as E2P | |
| import cv2 | |
| import numpy as np | |
| os.system("python setup.py build develop --user") | |
| os.system("pip install packaging==21.3") | |
| warnings.filterwarnings("ignore") | |
| from groundingdino.models import build_model | |
| from groundingdino.util.slconfig import SLConfig | |
| from groundingdino.util.utils import clean_state_dict | |
| from groundingdino.util.inference import annotate, load_image, predict | |
| import groundingdino.datasets.transforms as T | |
| from huggingface_hub import hf_hub_download | |
| picture_height = 480 | |
| picture_width = 720 | |
| picture_fov = 45 | |
| # Use this command for evaluate the GLIP-T model | |
| config_file = "groundingdino/config/GroundingDINO_SwinT_OGC.py" | |
| ckpt_repo_id = "ShilongLiu/GroundingDINO" | |
| ckpt_filenmae = "groundingdino_swint_ogc.pth" | |
| def detection(image): | |
| sub_images = process_panorama(image) | |
| predict_images = [] | |
| for sub_image in sub_images: | |
| predict_images.append(run_grounding(sub_image)) | |
| return predict_images | |
| def process_panorama(image): | |
| equ = E2P.Equirectangular(image) | |
| y_axis = 0 | |
| sub_images = [] | |
| while y_axis <= 0: | |
| z_axis = -150 | |
| while z_axis <= 90: | |
| img = equ.GetPerspective(picture_fov, z_axis, y_axis, picture_height, picture_width) | |
| # cv2.imwrite(f'{directory_name}_{z_axis}z.jpg', img) | |
| sub_images.append(img) | |
| z_axis += picture_fov | |
| y_axis += picture_fov | |
| return sub_images | |
| def load_model_hf(model_config_path, repo_id, filename, device='cpu'): | |
| args = SLConfig.fromfile(model_config_path) | |
| model = build_model(args) | |
| args.device = device | |
| cache_file = hf_hub_download(repo_id=repo_id, filename=filename) | |
| checkpoint = torch.load(cache_file, map_location='cpu') | |
| log = model.load_state_dict(clean_state_dict(checkpoint['model']), strict=False) | |
| print("Model loaded from {} \n => {}".format(cache_file, log)) | |
| _ = model.eval() | |
| return model | |
| def image_transform_grounding(init_image): | |
| transform = T.Compose([ | |
| T.RandomResize([800], max_size=1333), | |
| T.ToTensor(), | |
| T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| image, _ = transform(init_image, None) # 3, h, w | |
| return init_image, image | |
| def image_transform_grounding_for_vis(init_image): | |
| transform = T.Compose([ | |
| T.RandomResize([800], max_size=1333), | |
| ]) | |
| image, _ = transform(init_image, None) # 3, h, w | |
| return image | |
| model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae) | |
| def run_grounding(input_image): | |
| pil_img = Image.fromarray(input_image) | |
| init_image = pil_img.convert("RGB") | |
| grounding_caption = "traffic sign, car" | |
| box_threshold = 0.25 | |
| text_threshold = 0.25 | |
| _, image_tensor = image_transform_grounding(init_image) | |
| image_pil: Image = image_transform_grounding_for_vis(init_image) | |
| boxes, logits, phrases = predict(model, image_tensor, grounding_caption, box_threshold, text_threshold, | |
| device='cpu') | |
| annotated_frame = annotate(image_source=np.asarray(image_pil), boxes=boxes, logits=logits, phrases=phrases) | |
| image_with_box = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)) | |
| return image_with_box | |
| if __name__ == "__main__": | |
| detect_app = gr.Blocks() | |
| with detect_app: | |
| gr.Markdown("# Panorama Traffic Sign Detection Demo") | |
| gr.Markdown("Please upload an panorama picture to see the detection results.") | |
| gr.Markdown("Note the model runs on CPU for demo, so it may take a while to run the model.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(source='upload', type="numpy", label="Please upload a panorama picture.") | |
| run_button = gr.Button(value="Process & Detect", variant="primary") | |
| with gr.Row(): | |
| with gr.Column(): | |
| gallery = gr.Gallery(label="Detection Results").style( | |
| grid=(2, 3), preview=False, object_fit="none") | |
| run_button.click(fn=detection, inputs=[ | |
| input_image], outputs=[gallery]) | |
| detect_app.launch(share=False, show_api=False, show_error=True) | |