Spaces:
Runtime error
Runtime error
File size: 3,905 Bytes
94cd336 |
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 100 101 102 103 104 105 106 |
import torch
from efficientnet_pytorch import EfficientNet
from torchvision import transforms
from PIL import Image
import gradio as gr
from super_gradients.training import models
import cv2
import numpy as np
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the YOLO-NAS model
yolo_nas_l = models.get("yolo_nas_l", pretrained_weights="coco")
def bounding_boxes_overlap(box1, box2):
"""Check if two bounding boxes overlap or touch."""
x1, y1, x2, y2 = box1
x3, y3, x4, y4 = box2
return not (x3 > x2 or x4 < x1 or y3 > y2 or y4 < y1)
def merge_boxes(box1, box2):
"""Return the encompassing bounding box of two boxes."""
x1, y1, x2, y2 = box1
x3, y3, x4, y4 = box2
x = min(x1, x3)
y = min(y1, y3)
w = max(x2, x4)
h = max(y2, y4)
return (x, y, w, h)
def save_merged_boxes(predictions, image_np):
"""Save merged bounding boxes as separate images."""
processed_boxes = set()
roi = None # Initialize roi to None
for image_prediction in predictions:
bboxes = image_prediction.prediction.bboxes_xyxy
for box1 in bboxes:
for box2 in bboxes:
if np.array_equal(box1, box2):
continue
if bounding_boxes_overlap(box1, box2) and tuple(box1) not in processed_boxes and tuple(box2) not in processed_boxes:
merged_box = merge_boxes(box1, box2)
roi = image_np[int(merged_box[1]):int(merged_box[3]), int(merged_box[0]):int(merged_box[2])]
processed_boxes.add(tuple(box1))
processed_boxes.add(tuple(box2))
break # Exit the inner loop once a match is found
if roi is not None:
break # Exit the outer loop once a match is found
return roi
# Load the EfficientNet model
def load_model(model_path):
model = torch.load(model_path)
model = model.to(device)
model.eval() # Set the model to evaluation mode
return model
# Perform inference on an image
def predict_image(image, model):
# First, get the ROI using YOLO-NAS
image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
predictions = yolo_nas_l.predict(image_np, iou=0.3, conf=0.35)
roi_new = save_merged_boxes(predictions, image_np)
if roi_new is None:
roi_new = image_np # Use the original image if no ROI is found
# Convert ROI back to PIL Image for EfficientNet
roi_image = Image.fromarray(cv2.cvtColor(roi_new, cv2.COLOR_BGR2RGB))
# Define the image transformations
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# Convert PIL Image to Tensor
roi_image_tensor = transform(roi_image).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(roi_image_tensor)
_, predicted = outputs.max(1)
prediction_text = 'Accident' if predicted.item() == 0 else 'No accident'
return roi_image, prediction_text # Return both the roi_image and the prediction text
# Load the EfficientNet model outside the function to avoid loading it multiple times
model_path = 'vehicle.pt'
model = load_model(model_path)
# Gradio UI
title = "Vehicle Collision Classification"
description = "Upload an image to determine if it depicts a vehicle accident. Powered by EfficientNet."
examples = [["roi_none.png"], ["test2.jpeg"]] # Replace with your example image path
gr.Interface(fn=lambda img: predict_image(img, model),
inputs=gr.inputs.Image(type="pil"),
outputs=[gr.outputs.Image(type="pil"), "text"], # Updated outputs to handle both image and text
title=title,
description=description,
examples=examples).launch()
|