Mercurial > repos > bgruening > json2yolosegment

comparison yolov8.py @ 1:dfda27273ead draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 9685f843a52451b3416416094cc0e740f8825dcc
author bgruening
date Mon, 07 Jul 2025 06:47:08 +0000
parents 252fd085940d
children 158e6ce48345
comparison
equal deleted inserted replaced
0:252fd085940d 1:dfda27273ead
410 print(colored(f"\nYOLO prediction done in : '{elapsed}' sec \n", 'white', 'on_cyan')) 410 print(colored(f"\nYOLO prediction done in : '{elapsed}' sec \n", 'white', 'on_cyan'))
411 411
412 if (args.mode == "detect"): 412 if (args.mode == "detect"):
413 # Save bounding boxes 413 # Save bounding boxes
414 save_yolo_bounding_boxes_to_txt(predictions, args.save_dir) 414 save_yolo_bounding_boxes_to_txt(predictions, args.save_dir)
415
416 # Loop over each result
417 for result in predictions:
418 img = np.copy(result.orig_img)
419 image_filename = pathlib.Path(result.path).stem
420 overlay_path = os.path.join(args.save_dir, f"{image_filename}_overlay.jpg")
421
422 for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf):
423 x1, y1, x2, y2 = map(int, box.tolist())
424 class_num = int(cls.item())
425 confidence = conf.item()
426 label = f"{class_names[class_num]} {confidence:.2f}" if class_names else f"{class_num} {confidence:.2f}"
427
428 cv2.rectangle(img, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=2)
429 cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
430 0.5, (0, 255, 0), thickness=1)
431
432 cv2.imwrite(overlay_path, img)
433 print(colored(f"Overlay image saved at: {overlay_path}", 'cyan'))
415 elif (args.mode == "track"): 434 elif (args.mode == "track"):
416 results = model.track(source=datapath_for_prediction, 435 results = model.track(source=datapath_for_prediction,
417 tracker=args.tracker_file, 436 tracker=args.tracker_file,
418 conf=args.confidence, 437 conf=args.confidence,
419 iou=args.iou, 438 iou=args.iou,
449 print(colored(f"Tracking results saved in : '{args.save_dir}' \n", 'green')) 468 print(colored(f"Tracking results saved in : '{args.save_dir}' \n", 'green'))
450 elif (args.mode == "segment"): 469 elif (args.mode == "segment"):
451 # Read class names from the file 470 # Read class names from the file
452 with open(args.class_names_file, 'r') as f: 471 with open(args.class_names_file, 'r') as f:
453 class_names = [line.strip() for line in f.readlines()] 472 class_names = [line.strip() for line in f.readlines()]
454 # Create a mapping from class names to indices
455 class_to_index = {class_name: i for i, class_name in enumerate(class_names)} 473 class_to_index = {class_name: i for i, class_name in enumerate(class_names)}
456 474
457 # Save polygon coordinates 475 # Save polygon coordinates
458 for result in predictions: 476 for result in predictions:
459 # Create binary mask
460 img = np.copy(result.orig_img) 477 img = np.copy(result.orig_img)
461 filename = pathlib.Path(result.path).stem 478 filename = pathlib.Path(result.path).stem
462 b_mask = np.zeros(img.shape[:2], np.uint8) 479 b_mask = np.zeros(img.shape[:2], np.uint8)
463 mask_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + "_mask.tiff")).absolute()) 480 mask_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + "_mask.tiff")).absolute())
464 # Define output file path for text file
465 output_filename = os.path.splitext(filename)[0] + ".txt"
466 txt_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + ".txt")).absolute()) 481 txt_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + ".txt")).absolute())
467 482
468 for c, ci in enumerate(result): 483 for c, ci in enumerate(result):
469 # Extract contour result 484 if ci.masks is not None and ci.masks.xy:
470 contour = ci.masks.xy.pop() 485 # Extract contour
471 # Changing the type 486 contour = ci.masks.xy.pop()
472 contour = contour.astype(np.int32) 487 contour = contour.astype(np.int32).reshape(-1, 1, 2)
473 # Reshaping 488 _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)
474 contour = contour.reshape(-1, 1, 2) 489
475 # Draw contour onto mask 490 # Normalized polygon points
476 _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED) 491 points = ci.masks.xyn.pop()
477 492 obj_class = int(ci.boxes.cls.to("cpu").numpy().item())
478 # Normalized polygon points 493 confidence = result.boxes.conf.to("cpu").numpy()[c]
479 points = ci.masks.xyn.pop() 494
480 obj_class = int(ci.boxes.cls.to("cpu").numpy().item()) 495 with open(txt_save_as, 'a') as f:
481 confidence = result.boxes.conf.to("cpu").numpy()[c] 496 segmentation_points = ['{} {}'.format(points[i][0], points[i][1]) for i in range(len(points))]
482 497 segmentation_points_string = ' '.join(segmentation_points)
483 with open(txt_save_as, 'a') as f: 498 line = '{} {} {}\n'.format(obj_class, segmentation_points_string, confidence)
484 segmentation_points = ['{} {}'.format(points[i][0], points[i][1]) for i in range(len(points))] 499 f.write(line)
485 segmentation_points_string = ' '.join(segmentation_points) 500 else:
486 line = '{} {} {}\n'.format(obj_class, segmentation_points_string, confidence) 501 print(colored(f"⚠️ No mask found for object {c} in '{filename}'. Skipping.", "yellow"))
487 f.write(line) 502
488 503 # Overlay mask onto original image
489 imwrite(mask_save_as, b_mask, imagej=True) # save image 504 colored_mask = cv2.merge([b_mask, np.zeros_like(b_mask), np.zeros_like(b_mask)])
490 print(colored(f"Saved cropped image as : \n '{mask_save_as}' \n", 'magenta')) 505 blended = cv2.addWeighted(img, 1.0, colored_mask, 0.5, 0)
506 overlay_path = os.path.join(args.save_dir, filename + "_overlay.jpg")
507 cv2.imwrite(overlay_path, blended)
508
509 imwrite(mask_save_as, b_mask, imagej=True)
510 print(colored(f"Saved binary mask as : \n '{mask_save_as}' \n", 'magenta'))
491 print(colored(f"Polygon coordinates saved as : \n '{txt_save_as}' \n", 'cyan')) 511 print(colored(f"Polygon coordinates saved as : \n '{txt_save_as}' \n", 'cyan'))
492
493 else:
494 raise Exception(("Currently only 'detect' and 'segment' modes are available"))