annading commited on
Commit
775d75f
·
1 Parent(s): 553f53a

fixed rgb issue

Browse files
Files changed (1) hide show
  1. utils.py +7 -9
utils.py CHANGED
@@ -18,7 +18,7 @@ def plot_predictions(
18
  ) -> np.ndarray:
19
 
20
  image_source = cv2.imread(image)
21
- image_source = cv2.cvtColor(image_source, cv2.COLOR_BGR2RGB)
22
 
23
  boxes = sv.Detections(xyxy=boxes)
24
 
@@ -28,7 +28,7 @@ def plot_predictions(
28
  in zip(labels, scores)
29
  ]
30
 
31
- height, width, _ = image_source.shape
32
  thickness = math.ceil(width/200)
33
  text_scale = width/1500
34
  text_thickness = math.ceil(text_scale*1.5)
@@ -37,20 +37,18 @@ def plot_predictions(
37
  label_annotator = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX, text_scale=text_scale, text_thickness=text_thickness)
38
 
39
  # Create a semi-transparent overlay
40
- overlay = image_source.copy()
41
 
42
  # Apply bounding box annotations to the overlay
43
  overlay = bbox_annotator.annotate(scene=overlay, detections=boxes)
44
  overlay = label_annotator.annotate(scene=overlay, detections=boxes, labels=labels)
45
 
46
- # annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR)
47
- # annotated_frame = bbox_annotator.annotate(scene=annotated_frame, detections=boxes)
48
- # annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=boxes, labels=labels)
49
-
50
  # Blend overlay with original image using the specified opacity
51
- annotated_frame = cv2.addWeighted(overlay, opacity, image_source, 1 - opacity, 0)
 
 
52
 
53
- return annotated_frame
54
 
55
  def mp4_to_png(input_path: str, save_path: str, scale_factor: float) -> str:
56
  """ Converts mp4 to pngs for each frame of the video.
 
18
  ) -> np.ndarray:
19
 
20
  image_source = cv2.imread(image)
21
+ image_source_rgb = cv2.cvtColor(image_source, cv2.COLOR_BGR2RGB)
22
 
23
  boxes = sv.Detections(xyxy=boxes)
24
 
 
28
  in zip(labels, scores)
29
  ]
30
 
31
+ height, width, _ = image_source_rgb.shape
32
  thickness = math.ceil(width/200)
33
  text_scale = width/1500
34
  text_thickness = math.ceil(text_scale*1.5)
 
37
  label_annotator = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX, text_scale=text_scale, text_thickness=text_thickness)
38
 
39
  # Create a semi-transparent overlay
40
+ overlay = image_source_rgb.copy()
41
 
42
  # Apply bounding box annotations to the overlay
43
  overlay = bbox_annotator.annotate(scene=overlay, detections=boxes)
44
  overlay = label_annotator.annotate(scene=overlay, detections=boxes, labels=labels)
45
 
 
 
 
 
46
  # Blend overlay with original image using the specified opacity
47
+ annotated_frame = cv2.addWeighted(overlay, opacity, image_source_rgb, 1 - opacity, 0)
48
+
49
+ annotated_frame_bgr = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
50
 
51
+ return annotated_frame_bgr
52
 
53
  def mp4_to_png(input_path: str, save_path: str, scale_factor: float) -> str:
54
  """ Converts mp4 to pngs for each frame of the video.