Spaces:
Running
Running
File size: 3,636 Bytes
3ca5c3c cd299d7 c4cfa24 cd299d7 d558dbd cd299d7 3ca5c3c cd299d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | 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() |