Spaces:
Running
Running
| from pprint import pprint | |
| import torch | |
| import torchvision.models as tv_models | |
| import torchvision.transforms as transforms | |
| import torchvision.utils as vutils | |
| from IPython.display import Image as DisplayImage | |
| from PIL import Image | |
| from torchvision.io import decode_image | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| from torchvision.transforms.functional import to_pil_image | |
| from IPython.display import display | |
| import gradio as gr | |
| bb_model_weights = tv_models.detection.FasterRCNN_ResNet50_FPN_Weights.DEFAULT | |
| model = tv_models.detection.fasterrcnn_resnet50_fpn(weights=bb_model_weights).eval() | |
| def detect_and_draw_bboxes(model, image_path, object_indices, labels, bbox_colors, threshold, bbox_width=5): | |
| # Open and transform the image, and prepare the result tensor | |
| pil_image = Image.open(image_path).convert("RGB") | |
| transform_to_tensor = transforms.Compose([transforms.ToTensor()]) | |
| tensor_image_batch = transform_to_tensor(pil_image).unsqueeze(0) | |
| result_image_tensor = (tensor_image_batch.squeeze(0) * 255).byte() | |
| # Perform inference to get predictions for all possible objects | |
| with torch.no_grad(): | |
| prediction = model(tensor_image_batch)[0] | |
| # Initialize lists to collect all boxes, labels, and colors that meet the criteria | |
| all_boxes_to_draw = [] | |
| all_labels_to_draw = [] | |
| all_colors_to_draw = [] | |
| # Loop through each target class to find its boxes | |
| for index, label, color in zip(object_indices, labels, bbox_colors): | |
| # Filter predictions for the current class index and confidence threshold | |
| class_mask = (prediction['labels'] == index) & (prediction['scores'] > threshold) | |
| # Get the boxes for the current class | |
| boxes_for_this_class = prediction['boxes'][class_mask] | |
| if boxes_for_this_class.nelement() > 0: | |
| # Add the found boxes to our master list | |
| all_boxes_to_draw.extend(boxes_for_this_class.tolist()) | |
| # Create and add corresponding labels and colors | |
| all_labels_to_draw.extend([label] * len(boxes_for_this_class)) | |
| all_colors_to_draw.extend([color] * len(boxes_for_this_class)) | |
| # After checking all classes, draw all collected boxes at once if any were found | |
| if all_boxes_to_draw: | |
| result_image_tensor = vutils.draw_bounding_boxes( | |
| result_image_tensor, | |
| torch.tensor(all_boxes_to_draw), | |
| labels=all_labels_to_draw, | |
| colors=all_colors_to_draw, | |
| width=bbox_width | |
| ) | |
| else: | |
| # If the list of boxes to draw is empty, print this information. | |
| print(f"No objects from the list {labels} were found with a confidence score above {threshold}.\n") | |
| return result_image_tensor | |
| def run_detection(image_path): | |
| result_image_tensor = detect_and_draw_bboxes( | |
| model=model, | |
| image_path=image_path, | |
| object_indices=[3, 10], | |
| labels=['car', 'traffic light'], | |
| bbox_colors=['red', 'blue'], | |
| threshold=0.7, | |
| ) | |
| return to_pil_image(result_image_tensor) | |
| with gr.Blocks() as demo: | |
| iface = gr.Interface( | |
| fn=run_detection, | |
| inputs=gr.Image(type="filepath", label="Upload Image"), | |
| outputs=gr.Image(type="pil", label="Output Image"), | |
| title="<br>This app uses:<br>PyTorch Vision <br>PyTorch Model: FasterRCNN_ResNet50_FPN<br>PS: click submit after image uploaded" | |
| ) | |
| gr.Markdown("### The model is adjusted to only detect car and traffic light. Example output:") | |
| gr.Image( | |
| value="exap1.JPG", | |
| show_label=False, | |
| width=450 | |
| ) | |
| demo.launch() |