draw.py 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. import matplotlib.pyplot as plt
  2. import pydicom
  3. import cv2
  4. import numpy as np
  5. from skimage.measure import find_contours
  6. # def get_mask_contour(mask):
  7. # mskray = mask.astype(np.uint8)
  8. # edged_mask = cv2.Canny(mskray, np.min(mskray), np.max(mskray))
  9. # return edged_mask
  10. # Global variables
  11. drawing = False # True when the mouse button is pressed
  12. points = [] # Stores points of the ROI
  13. mask = None # To store the final mask
  14. def dcm2cv (dicom_img):
  15. # Normalize the pixel values to the range 0-255
  16. image_normalized = cv2.normalize(dicom_img, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
  17. # Convert to unsigned 8-bit integer type
  18. image_normalized = image_normalized.astype(np.uint8)
  19. cv_image = cv2.merge([image_normalized, image_normalized, image_normalized])
  20. return cv_image
  21. # Mouse callback function
  22. def draw_roi(event, x, y, flags, param):
  23. global drawing, points, mask
  24. if event == cv2.EVENT_LBUTTONDOWN: # Start drawing
  25. drawing = True
  26. points.append((x, y))
  27. elif event == cv2.EVENT_MOUSEMOVE and drawing: # Add points while moving
  28. points.append((x, y))
  29. elif event == cv2.EVENT_LBUTTONUP: # Stop drawing
  30. drawing = False
  31. elif event == cv2.EVENT_RBUTTONDOWN: # Right-click to complete the ROI
  32. print('Drawing Complete')
  33. if len(points) > 2:
  34. mask = np.zeros_like(image, dtype=np.uint8)
  35. pts = np.array(points, np.int32).reshape((-1, 1, 2))
  36. cv2.fillPoly(mask, [pts], (255, 255, 255))
  37. contours = find_contours(mask[:, :, 0], level=0.5)
  38. # plt.imshow(image)
  39. # for cn in range(len(contours)):
  40. # plt.plot(contours[cn][:, 1], contours[cn][:, 0], 'r')
  41. # plt.show()
  42. # Load an image (replace this with actual MR image loading)
  43. path = r'E:\projects\knee_seg\data_tamplet\PA0\ST0\SE4\IM40'
  44. file_dcm = pydicom.read_file(path)
  45. image = file_dcm.pixel_array
  46. image = dcm2cv(image)
  47. cv2.namedWindow("Image")
  48. cv2.setMouseCallback("Image", draw_roi)
  49. while True:
  50. temp_img = image.copy()
  51. print(temp_img.shape)
  52. if points:
  53. cv2.polylines(temp_img, [np.array(points)], isClosed=False, color=(255, 0, 0), thickness=2)
  54. cv2.imshow("Image", temp_img)
  55. key = cv2.waitKey(1) & 0xFF
  56. if key == 13 and mask is not None: # Press 'enter' to save the mask
  57. cv2.imwrite("mask.png", mask)
  58. print("Mask saved as mask.png")
  59. break
  60. elif key == 27: # Press 'Esc' to exit
  61. # print(mask)
  62. cv2.imwrite("mask.png", mask)
  63. print("Break: Mask saved ")
  64. break
  65. cv2.destroyAllWindows()