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