import torch import cv2 import numpy as np from pathlib import Path from PIL import Image, ImageDraw from models.common import DetectMultiBackend from utils.augmentations import letterbox from utils.general import non_max_suppression, scale_boxes from models.experimental import attempt_load from utils.torch_utils import select_device, smart_inference_mode from tqdm import tqdm #################################################### #################### Parameters #################### # Define your object class indices and corresponding colors ## 사람 (Person): 0 ## 자전거 (Bicycle): 1 ## 자동차 (Car): 2 ## 오토바이 (Motorcycle): 3 ## 버스 (Bus): 5 ## 트럭 (Truck): 7 ## 우산 (Umbrella): 25 obj_class = [0, 1, 2, 3, 5, 7, 25] # class별 가중치값 class_score_weights = [1, 3, 12, 3, 20, 20, 0.3] # 경고, 위험 criterion = [50,70] criterion_color = [(255,255,255),(0,255,255),(0,0,255)] # 밤시간 패널티 night_time = False night_penalty = 0.8 # Define zones zone1 = [(300, 330), (448, 330), (849, 710), (329, 710)] zone2 = [(448, 330), (620, 330), (1273, 710), (849, 710)] zones = [zone1, zone2] # If you have more zones, add them to this list # input video path vid_path = 'daytime.mp4' # object detection threshold conf_thres=0.25 # confidence threshold iou_thres=0.45 # NMS IOU threshold #################################################### #################################################### # if night then apply night time penalty if night_time: criterion = criterion*night_penalty # Mapping of class indices to their weights class_weights = {cls: weight for cls, weight in zip(obj_class, class_score_weights)} # Update the draw_polygon function to accept a color parameter def draw_polygon(frame, danger_points, color): img = Image.fromarray(frame) draw = ImageDraw.Draw(img) draw.polygon(danger_points, outline=color, width=5) return np.array(img) # check if the center bottom of the bounding box is inside the polygon def is_point_inside_polygon(x, y, polygon): n = len(polygon) inside = False p1x, p1y = polygon[0] for i in range(n+1): p2x, p2y = polygon[i % n] if y > min(p1y, p2y): if y <= max(p1y, p2y): if x <= max(p1x, p2x): if p1y != p2y: xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x if p1x == p2x or x <= xinters: inside = not inside p1x, p1y = p2x, p2y return inside # output video path out_path = vid_path.split('.')[0]+'_result.mp4' # Model weights = './yolov5x.pt' device = select_device(0) model = DetectMultiBackend(weights, device=device, dnn=False, data='data/coco128.yaml', fp16=False) stride = model.stride # model stride names = model.module.names if hasattr(model, 'module') else model.names # capture video cap = cv2.VideoCapture(vid_path) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = int(cap.get(cv2.CAP_PROP_FPS)) # output video writer fourcc = cv2.VideoWriter_fourcc(*'MP4V') out = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) # get the total number of frames in the video total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) pbar = tqdm(total=total_frames, desc="Processing frames", unit="frame") while True: ret, frame = cap.read() if not ret: break # Resize and pad image while meeting stride-multiple constraints img = letterbox(frame, new_shape=640, stride=stride)[0] # Convert img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3xHxW img = np.ascontiguousarray(img) # Convert to torch tensor img = torch.from_numpy(img).to(device) img = img.float() # uint8 to fp32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference pred = model(img, augment=False)[0] # Apply NMS pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False) # Initialize weighted scores for each zone zone_scores = [0] * len(zones) dets = [] # Process detections for i, det in enumerate(pred): # detections per image if len(det): # Rescale boxes from img_size to original size det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], frame.shape).round() # Process detections for *xyxy, conf, cls in det: cls_int = int(cls) if cls_int in obj_class: # Get bounding box coordinates x1 = int(xyxy[0].cpu().numpy()) y1 = int(xyxy[1].cpu().numpy()) x2 = int(xyxy[2].cpu().numpy()) y2 = int(xyxy[3].cpu().numpy()) # Calculate the center bottom of the bounding box x_center = int((x1 + x2) / 2) y_center = int(y2) # Check which zone the detection is in and update the score for zone_idx, zone in enumerate(zones): if is_point_inside_polygon(x_center, y_center, zone): # Add the weighted score for this class to the zone's score zone_scores[zone_idx] += class_weights[cls_int] break # Draw zones on the frame with updated colors based on the score for zone_idx, score in enumerate(zone_scores): # Determine color based on score if score < criterion[0]: color = criterion_color[0] elif score < criterion[1]: color = criterion_color[1] else: color = criterion_color[2] # Now draw the zone with the determined color frame = draw_polygon(frame, zones[zone_idx], color) # Display scores on the frame with the same color as the zone cv2.putText(frame, f"Zone {zone_idx+1}: {score:.2f}", (10, 30 + zone_idx*30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2, cv2.LINE_AA) # Write frame to output video out.write(frame) pbar.update(1) # Display the resulting frame cv2.imshow('Frame', frame) # Wait for the 'q' key to be pressed to exit if cv2.waitKey(1) & 0xFF == ord('q'): break # release video capture and writer cap.release() out.release() pbar.close()