Skip to content

Commit 7bbeea0

Browse files
Koldim2001Koldim2001
committed
fix params
1 parent 5eb33e0 commit 7bbeea0

3 files changed

Lines changed: 166 additions & 12 deletions

File tree

‎README.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -111,10 +111,10 @@ Class implementing cropping and passing crops through a neural network for detec
111111
| **Argument** | **Type** | **Default** | **Description** |
112112
|-----------------------|------------------------|--------------|----------------------------------------------------------------------------------------------------------------|
113113
| image | np.ndarray | | The input image in BGR format. |
114-
| model_path | str | "yolov8m.pt" | Path to the YOLO model. |
114+
| model_path | str | "yolo11m.pt" | Path to the YOLO model. |
115115
| model | ultralytics model | None | Pre-initialized model object. If provided, the model will be used directly instead of loading from model_path. |
116116
| imgsz | int | 640 | Size of the input image for inference YOLO. |
117-
| conf | float | 0.5 | Confidence threshold for detections YOLO. |
117+
| conf | float | 0.25 | Confidence threshold for detections YOLO. |
118118
| iou | float | 0.7 | IoU threshold for non-maximum suppression YOLOv8 of single crop. |
119119
| classes_list | List[int] or None | None | List of classes to filter detections. If None, all classes are considered. |
120120
| segment | bool | False | Whether to perform segmentation (if the model supports it). |
@@ -167,7 +167,7 @@ Possible arguments of the ```visualize_results``` function:
167167
| thickness | int | 4 | The thickness of bounding box and text. |
168168
| font | cv2.font |cv2.FONT_HERSHEY_SIMPLEX | The font type for class labels. |
169169
| font_scale | float | 1.5 | The scale factor for font size. |
170-
| delta_colors | int | seed=0 | The random seed offset for color variation. |
170+
| delta_colors | int | seed=3 | The random seed offset for color variation. |
171171
| dpi | int | 150 | Final visualization size (plot is bigger when dpi is higher). |
172172
| random_object_colors | bool | False | If true, colors for each object are selected randomly. |
173173
| show_confidences | bool | False | If true and show_class=True, confidences near class are visualized. |

‎patched_yolo_infer/functions_extra.py‎

Lines changed: 161 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ def visualize_results_usual_yolo_inference(
1010
img,
1111
model,
1212
imgsz=640,
13-
conf=0.5,
13+
conf=0.25,
1414
iou=0.7,
1515
segment=False,
1616
show_boxes=True,
@@ -22,7 +22,7 @@ def visualize_results_usual_yolo_inference(
2222
thickness=4,
2323
font=cv2.FONT_HERSHEY_SIMPLEX,
2424
font_scale=1.5,
25-
delta_colors=0,
25+
delta_colors=3,
2626
dpi=150,
2727
random_object_colors=False,
2828
show_confidences=False,
@@ -33,13 +33,13 @@ def visualize_results_usual_yolo_inference(
3333
inference_extra_args=None,
3434
):
3535
"""
36-
Visualizes the results of usual YOLOv8 or YOLOv8-seg inference on an image
36+
Visualizes the results of usual YOLO or YOLO-seg inference on an image
3737
3838
Args:
3939
img (numpy.ndarray): The input image in BGR format.
4040
model: The object detection or segmentation model (yolov8).
4141
imgsz (int): The input image size for the model. Default is 640.
42-
conf (float): The confidence threshold for detection. Default is 0.5.
42+
conf (float): The confidence threshold for detection. Default is 0.25.
4343
iou (float): The intersection over union threshold for detection. Default is 0.7.
4444
segment (bool): Whether to perform instance segmentation. Default is False.
4545
show_boxes (bool): Whether to show bounding boxes. Default is True.
@@ -51,7 +51,7 @@ def visualize_results_usual_yolo_inference(
5151
thickness (int): The thickness of bounding box and text. Default is 4.
5252
font: The font type for class labels. Default is cv2.FONT_HERSHEY_SIMPLEX.
5353
font_scale (float): The scale factor for font size. Default is 1.5.
54-
delta_colors (int): The random seed offset for color variation. Default is 0.
54+
delta_colors (int): The random seed offset for color variation. Default is 3.
5555
dpi (int): Final visualization size (plot is bigger when dpi is higher).
5656
random_object_colors (bool): If True, colors for each object are selected randomly.
5757
show_confidences (bool): If True and show_class=True, confidences near class are visualized.
@@ -177,6 +177,160 @@ def visualize_results_usual_yolo_inference(
177177
plt.show()
178178

179179

180+
def visualize_results_yolo_pose_inference(
181+
img,
182+
model,
183+
imgsz=640,
184+
conf=0.25,
185+
iou=0.7,
186+
show_boxes=True,
187+
show_class=True,
188+
color_class_background=(0, 0, 255),
189+
color_class_text=(255, 255, 255),
190+
thickness=4,
191+
point_radius=4,
192+
connection_schema=None,
193+
min_landmark_visibility=0.25,
194+
font=cv2.FONT_HERSHEY_SIMPLEX,
195+
font_scale=1.5,
196+
delta_colors=3,
197+
dpi=150,
198+
random_object_colors=False,
199+
show_confidences=False,
200+
axis_off=True,
201+
show_classes_list=[],
202+
list_of_class_colors=None,
203+
return_image_array=False,
204+
inference_extra_args=None,
205+
):
206+
"""
207+
Visualizes the results of usual YOLO-pose inference on an image
208+
209+
Args:
210+
img (numpy.ndarray): The input image in BGR format.
211+
model: The object detection or segmentation model (yolov8).
212+
imgsz (int): The input image size for the model. Default is 640.
213+
conf (float): The confidence threshold for detection. Default is 0.25.
214+
iou (float): The intersection over union threshold for detection. Default is 0.7.
215+
show_boxes (bool): Whether to show bounding boxes. Default is True.
216+
show_class (bool): Whether to show class labels. Default is True.
217+
color_class_background (tuple / list of tuple): The background BGR color for class labels. Default is (0, 0, 255) (red).
218+
color_class_text (tuple): The text BGR color for class labels. Default is (255, 255, 255) (white).
219+
thickness (int): The thickness of bounding box and text. Default is 4.
220+
point_radius (int): The radius of the landmark points to be drawn on the image.
221+
connection_schema (list): A list of tuples defining how landmarks should be connected to form a skeleton.
222+
Each tuple contains two indices representing the landmarks to be connected.
223+
If None or empty, only landmarks will be drawn without any connections.
224+
min_landmark_visibility (float): The minimum confidence threshold for a landmark's visibility to be drawn.
225+
font: The font type for class labels. Default is cv2.FONT_HERSHEY_SIMPLEX.
226+
font_scale (float): The scale factor for font size. Default is 1.5.
227+
delta_colors (int): The random seed offset for color variation. Default is 3.
228+
dpi (int): Final visualization size (plot is bigger when dpi is higher).
229+
random_object_colors (bool): If True, colors for each object are selected randomly.
230+
show_confidences (bool): If True and show_class=True, confidences near class are visualized.
231+
axis_off (bool): If True, axis is turned off in the final visualization.
232+
show_classes_list (list): If empty, visualize all classes. Otherwise, visualize only classes in the list.
233+
inference_extra_args (dict/None): Dictionary with extra ultralytics inference parameters.
234+
list_of_class_colors (list/None): A list of tuples representing the colors for each class in BGR format.
235+
If provided, these colors will be used for displaying the classes instead of random colors.
236+
The number of tuples in the list must match the number of possible classes in the network.
237+
return_image_array (bool): If True, the function returns the image bgr array instead of displaying it.
238+
Default is False.
239+
240+
Returns:
241+
None/np.array
242+
"""
243+
244+
# Perform inference
245+
extra_args = {} if inference_extra_args is None else inference_extra_args
246+
predictions = model.predict(img, imgsz=imgsz, conf=conf, iou=iou, verbose=False, **extra_args)
247+
248+
labeled_image = img.copy()
249+
250+
if random_object_colors:
251+
random.seed(int(delta_colors))
252+
253+
class_names = model.names
254+
255+
# Process each prediction
256+
for pred in predictions:
257+
258+
# Get the bounding boxes and convert them to a list of lists
259+
boxes = pred.boxes.xyxy.cpu().int().tolist()
260+
261+
# Get the classes and convert them to a list
262+
classes = pred.boxes.cls.cpu().int().tolist()
263+
264+
# Get the mask confidence scores
265+
confidences = pred.boxes.conf.cpu().numpy()
266+
267+
num_objects = len(classes)
268+
269+
# Visualization
270+
for i in range(num_objects):
271+
# Get the class for the current detection
272+
class_index = int(classes[i])
273+
class_name = class_names[class_index]
274+
275+
if show_classes_list and class_index not in show_classes_list:
276+
continue
277+
278+
if random_object_colors:
279+
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
280+
elif list_of_class_colors is None:
281+
# Assign color according to class
282+
random.seed(int(classes[i] + delta_colors))
283+
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
284+
else:
285+
color = list_of_class_colors[class_index]
286+
287+
box = boxes[i]
288+
x_min, y_min, x_max, y_max = box
289+
290+
# Write class label
291+
if show_boxes:
292+
cv2.rectangle(labeled_image, (x_min, y_min), (x_max, y_max), color, thickness)
293+
294+
if show_class:
295+
if show_confidences:
296+
label = f'{str(class_name)} {confidences[i]:.2}'
297+
else:
298+
label = str(class_name)
299+
(text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)
300+
background_color = (
301+
color_class_background[class_index]
302+
if isinstance(color_class_background, list)
303+
else color_class_background
304+
)
305+
cv2.rectangle(
306+
labeled_image,
307+
(x_min, y_min),
308+
(x_min + text_width + 5, y_min + text_height + 5),
309+
background_color,
310+
-1,
311+
)
312+
cv2.putText(
313+
labeled_image,
314+
label,
315+
(x_min + 5, y_min + text_height),
316+
font,
317+
font_scale,
318+
color_class_text,
319+
thickness=thickness,
320+
)
321+
322+
if return_image_array:
323+
return labeled_image
324+
else:
325+
# Display the final image with overlaid masks and labels
326+
plt.figure(figsize=(8, 8), dpi=dpi)
327+
labeled_image = cv2.cvtColor(labeled_image, cv2.COLOR_BGR2RGB)
328+
plt.imshow(labeled_image)
329+
if axis_off:
330+
plt.axis('off')
331+
plt.show()
332+
333+
180334
def get_crops(
181335
image_full,
182336
shape_x: int,
@@ -284,7 +438,7 @@ def visualize_results(
284438
thickness=4,
285439
font=cv2.FONT_HERSHEY_SIMPLEX,
286440
font_scale=1.5,
287-
delta_colors=0,
441+
delta_colors=3,
288442
dpi=150,
289443
random_object_colors=False,
290444
show_confidences=False,
@@ -313,7 +467,7 @@ def visualize_results(
313467
thickness (int): The thickness of bounding box and text. Default is 4.
314468
font: The font type for class labels. Default is cv2.FONT_HERSHEY_SIMPLEX.
315469
font_scale (float): The scale factor for font size. Default is 1.5.
316-
delta_colors (int): The random seed offset for color variation. Default is 0.
470+
delta_colors (int): The random seed offset for color variation. Default is 3.
317471
dpi (int): Final visualization size (plot is bigger when dpi is higher). Default is 150.
318472
random_object_colors (bool): If true, colors for each object are selected randomly. Default is False.
319473
show_confidences (bool): If true and show_class=True, confidences near class are visualized. Default is False.

‎patched_yolo_infer/nodes/MakeCropsDetectThem.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,9 +61,9 @@ class MakeCropsDetectThem:
6161
def __init__(
6262
self,
6363
image: np.ndarray,
64-
model_path="yolov8m.pt",
64+
model_path="yolo11m.pt",
6565
imgsz=640,
66-
conf=0.5,
66+
conf=0.25,
6767
iou=0.7,
6868
classes_list=None,
6969
segment=False,

0 commit comments

Comments
 (0)