crop_arena_pictures.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## README\n",
    "\n",
    "This script crops the arena pictures of each mouse track experiment.\n",
    "\n",
    "It is advisable to preprocess (e.g., in Photoshop) the arena pictures that will be processed by this script.\n",
    "- The original arena picture must keep all the landmarks but have its B&W contrast enhanced\n",
    "- A copy of the original arena picture with digitally removed landmarks is advisable to make hole identification correct.\n",
    "- Adjust the correction pixels in the tuple list of picture files below (last 3 values) to get better alignment results.\n",
    "\n",
    "The hole detection may detect multiple times a single hole, you must correct for that in the generated txt file by checking the hole index in the _holes.png gerenated picture\n",
    "Otherwise, you can digitally paint the faulty holes in photoshop before processing the arena picture with this script\n",
    "\n",
    "When files with _entrance.png are provided (entrance marks in black on a completely white background; obtained from the original picture),\n",
    "then the entrance coordinates are also written in the _holes.txt file using the special hole indices -4 to -1\n",
    "\n",
    "The idea here is to obtain an arena picture that is compatible with the arena picture of the first experiments (the \"pilot\" pictures in the main arena_picture folder)\n",
    "\n",
    "Outputs:\n",
    "    * a cropped arena picture with an img_extent\n",
    "        this file should be included in the process_mouse_trial_lib function \"get_arena_picture_bounds\"\n",
    "        with the proper \"if\" condition for that experiment (which is usually copied from Kelly's plot script, e.g. the experimemnt date most of the times)\n",
    "    * the img_extent for the cropped picture (printed on the screen)\n",
    "    * a BW arena picture with the (0,0)_world coordinates marked in blue\n",
    "    * a BW arena picture with all the detected holes identified by unique sequential indices\n",
    "    * a txt file with the (x,y)_world coordinates of each identified hole that is marked in the holes picture\n",
    "\n",
    "The (x0,x1)_world and (y0,y1)_world parameters in the list below are also copied from Kelly's plot script, and represent the original extent (or x,y axes bounds) of each full arena picture.\n",
    "\n",
    "\n",
    "\n",
    "## #################################################\n",
    "## #################################################\n",
    "## #################################################\n",
    "## #################################################\n",
    "\n",
    "## TO DO\n",
    "update the algorithm such that it reads some \"*_marks.png\" arena files, where the marks contain each object in a different color\n",
    "e.g.: arena boundarys in blue; entrance in red, green, magenta and cyan; targets in subtones of these colors corresponding to each entrance, holes color, etc...\n",
    "then, I can use the cluster-finding code to seek for the cluster of each specific color, and I don't need to use the circle detection algorightm\n",
    "it could have some input like this:\n",
    "'arena_picture/dec-2019/BKGDimage-localCues.png',                      'arena_picture/dec-2019/BKGDimage-localCues_preproc.png',        misc.structtype(arena=(0,0,255),holes=(0,0,0),entrances={0:(255,0,0),1:(0,255,0),2:(255,0,255),3:(0,255,255)},targets={0:(150,0,0),1:(0,150,0),2:(150,0,150),3:(0,150,150)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy\n",
    "import modules.io as io\n",
    "import modules.edge_detector as edged\n",
    "import modules.helper_func_class as misc\n",
    "import modules.process_mouse_trials_lib as plib\n",
    "import PIL\n",
    "import PIL.ImageDraw\n",
    "import PIL.ImageEnhance\n",
    "\n",
    "arena_pic_dir         = 'arena_picture'\n",
    "#arena_pic_dir          = 'debug_arena_picture'\n",
    "arena_pic_ref          = os.path.join(arena_pic_dir,'BKGDimage-pilot.png') #os.path.join(arena_pic_dir,plib.get_arena_all_picture_filename()['SW'])\n",
    "arena_ref_img_mark     = os.path.join(arena_pic_dir,'BKGDimage-pilot_mark.png')\n",
    "arena_pic_bounds       = misc.structtype(arena_pic_left=-97.0,arena_pic_right=97.0, arena_pic_bottom=-73.0,arena_pic_top=73.0) #plib.get_arena_picture_bounds()\n",
    "arena_radius_tolerance = 0.02 # tolerance for the apparent arena radius in a picture (if ratio of radius/radius_ref > tolerance, means the the picture is zoomed in/out in relation to the reference picture)\n",
    "\n",
    "delete_before_processing = True\n",
    "\n",
    "to_gray = lambda I: I.convert('LA') if 'A' in I.getbands() else I.convert('L')\n",
    "\n",
    "dx_ref = numpy.asarray((arena_pic_bounds.arena_pic_left   , arena_pic_bounds.arena_pic_right ))\n",
    "dy_ref = numpy.asarray((arena_pic_bounds.arena_pic_bottom ,  arena_pic_bounds.arena_pic_top  ))\n",
    "\n",
    "\"\"\"\n",
    "the mark_scheme varible represents the colors of each feature to be detected\n",
    "\n",
    "a file with suffix *_mark.png should be given for the arena picture to be processed according to the mark_scheme;\n",
    "otherwise, mark_scheme will NOT be used.\n",
    "\"\"\"\n",
    "mark_scheme = misc.structtype(arena_boundary='#0000FF', holes='#000000', entrances='#00FF00', targets='#FF0000')\n",
    "\n",
    "\"\"\"\n",
    "The tuple list below contains the following information in each tuple\n",
    "\n",
    "[ tuple_1, tuple_2, ... ]\n",
    "\n",
    "such that tuple_X has:\n",
    "\n",
    "   do_processing,                            # flag to determine whether this file will be processed\n",
    "   arena_file_name_inside_arena_picture_dir, # relative to the arena_picture directory, this is the file that will be rescaled and used to plot trajectories (must keep landmarks)\n",
    "   arena_preprocessed_file_name_if_needed,   # (optional, recommended) relative to the arena_picture directory, optional, this file is only used for hole detection, should have landmarks removed (digitally)\n",
    "   arena_entrance_file,                      # only entrance marks in white background\n",
    "   arena_marks_file,                         # (optional, recommended; makes the code faster) marks file with marks for entrance, holes, arena circle, and targets, respecting the mark_scheme color scheme\n",
    "   (x0,x1)_world,                            # the left and right boundaries (in cm) of the original arena picture given\n",
    "   (y0,y1)_world,                            # the bottom and top boundaries (in cm) of the original arena picture given\n",
    "   contrast_factor,                          # contrast enhancing factor (adjust at will to get better automatic arena and whole detection)\n",
    "   brightness_factor,                        # brightness enhancing factor (adjust at will to get better automatic arena and whole detection)\n",
    "   dx_px_correction,                         # correction factors for the corresponding coordinates in the arena and hole detection (in pixels)\n",
    "   dy_px_correction,                         # correction factors for the corresponding coordinates in the arena and hole detection (in pixels)\n",
    "   r_correction                              # correction factors for the corresponding coordinates in the arena and hole detection (in pixels)\n",
    "\"\"\"\n",
    "arena_pic = [\n",
    "   (False, 'dec-2019/BKGDimage-localCues.png'                      , 'dec-2019/BKGDimage-localCues_preproc.png'                      , 'dec-2019/BKGDimage-localCues_entrance.png'                      , ''                                                          , numpy.array((-129.58, 129.58)), numpy.array((-73.03, 73.03)),100.0,1.1, 0, 0,-2),\n",
    "   (False, 'mar-may-2021/BKGDimage-localCues_clear.png'            , 'mar-may-2021/BKGDimage-localCues_clear_preproc.png'            , 'mar-may-2021/BKGDimage-localCues_clear_entrance.png'            , ''                                                          , numpy.array((-151.06, 150.43)), numpy.array((-84.23, 84.86)),100.0,1.1, 0, 0,-4),\n",
    "   (False, 'mar-may-2021/BKGDimage-localCues_Letter.png'           , 'mar-may-2021/BKGDimage-localCues_Letter_preproc.png'           , 'mar-may-2021/BKGDimage-localCues_Letter_entrance.png'           , ''                                                          , numpy.array((-151.06, 150.43)), numpy.array((-84.23, 84.86)),100.0,1.1, 0, 0,-4),\n",
    "   (False, 'mar-may-2021/BKGDimage-localCues_LetterTex.png'        , 'mar-may-2021/BKGDimage-localCues_LetterTex_preproc.png'        , 'mar-may-2021/BKGDimage-localCues_LetterTex_entrance.png'        , ''                                                          , numpy.array((-151.06, 150.43)), numpy.array((-84.23, 84.86)),100.0,1.1, 0, 0,-4),\n",
    "   (False, 'jun-jul-aug-nov-2021/BKGDimage_3LocalCues.png'         , 'jun-jul-aug-nov-2021/BKGDimage_3LocalCues_preproc.png'         , 'jun-jul-aug-nov-2021/BKGDimage_3LocalCues_entrance.png'         , ''                                                          , numpy.array((-151.26, 150.94)), numpy.array((-84.66, 85.29)),100.0,1.1, 0, 0,-4),\n",
    "   (False, 'jun-jul-aug-nov-2021/BKGDimage_3LocalCues_Reverse.png' , 'jun-jul-aug-nov-2021/BKGDimage_3LocalCues_Reverse_preproc.png' , 'jun-jul-aug-nov-2021/BKGDimage_3LocalCues_Reverse_entrance.png' , ''                                                          , numpy.array((-151.26, 150.94)), numpy.array((-84.66, 85.29)),100.0,1.1, 0, 0,-4),\n",
    "   (False, 'jun-jul-aug-nov-2021/BKGDimage_asymm.png'              , 'jun-jul-aug-nov-2021/BKGDimage_asymm_preproc.png'              , 'jun-jul-aug-nov-2021/BKGDimage_asymm_entrance.png'              , ''                                                          , numpy.array((-151.26, 150.94)), numpy.array((-84.66, 85.29)),100.0,1.1, 0, 0,-2),\n",
    "   (False, 'jun-jul-aug-nov-2021/BKGDimage_asymmRev.png'           , 'jun-jul-aug-nov-2021/BKGDimage_asymmRev_preproc.png'           , 'jun-jul-aug-nov-2021/BKGDimage_asymmRev_entrance.png'           , ''                                                          , numpy.array((-151.26, 150.94)), numpy.array((-84.66, 85.29)),100.0,1.1, 0, 0,-2),\n",
    "   (False, 'jun-jul-aug-nov-2021/BKGDimage-arenaB.png'             , ''                                                              , 'jun-jul-aug-nov-2021/BKGDimage-arenaB_entrance.png'             , 'jun-jul-aug-nov-2021/BKGDimage-arenaB_mark.png'            , numpy.array((-151.26, 150.94)), numpy.array((-84.66, 85.29)),100.0,1.1,-2,-4,-5),\n",
    "   (False, 'jun-jul-aug-nov-2021/BDGDimage_arenaB_visualCues.png'  , ''                                                              , 'jun-jul-aug-nov-2021/BDGDimage_arenaB_visualCues_entrance.png'  , 'jun-jul-aug-nov-2021/BDGDimage_arenaB_visualCues_mark.png' , numpy.array((-151.26, 150.94)), numpy.array((-84.66, 85.29)),100.0,1.1,-2,-4,-2),\n",
    "   (False, 'jun-jul-aug-nov-2021/BKGDimage-Nov15.png'              , ''                                                              , 'jun-jul-aug-nov-2021/BKGDimage-Nov15_entrance.png'              , 'jun-jul-aug-nov-2021/BKGDimage-Nov15_mark.png'             , numpy.array((-151.26, 150.94)), numpy.array((-84.66, 85.29)),100.0,1.1,-2,-4,-2),\n",
    "   (False, '2022/BKGDimage-220812.png'                             , ''                                                              , '2022/BKGDimage-220812_entrance.png'                             , '2022/BKGDimage-220812_mark.png'                            , numpy.array((-151.65, 151.34)), numpy.array((-84.88, 85.51)),100.0,1.1, 0, 0, 0),\n",
    "   (False, '2022/BKGDimage-20221011.png'                           , ''                                                              , '2022/BKGDimage-20221011_entrance.png'                           , '2022/BKGDimage-20221011_mark.png'                          , numpy.array((-136.92, 136.63)), numpy.array((-76.78, 77.20)),100.0,1.1, 0, 0, 0),\n",
    "   (True , '2022/BKGDimage-20220920.png'                           , ''                                                              , ''                                                               , '2022/BKGDimage-20220920_mark.png'                          , numpy.array(( -99.34, 98.55 )), numpy.array((-84.36, 85.47)),100.0,1.1, 0, 0, 0),\n",
    "   (True , '2022/BKGDimage-20220920-probe2.png'                    , ''                                                              , ''                                                               , '2022/BKGDimage-20220920-probe2_mark.png'                   , numpy.array(( -99.34, 98.55 )), numpy.array((-84.36, 85.47)),100.0,1.1, 0, 0, 0),\n",
    "   (True , '2022/BKGDimage-20221104.png'                           , ''                                                              , ''                                                               , '2022/BKGDimage-20221104_mark.png'                          , numpy.array((-151.65, 151.34)), numpy.array((-84.88, 85.51)),100.0,1.1, 0, 0, 0)\n",
    "   ]\n",
    "\n",
    "\n",
    "mark_pic_not_required = [\n",
    "   'dec-2019/BKGDimage-localCues.png'                     ,\n",
    "   'mar-may-2021/BKGDimage-localCues_clear.png'           ,\n",
    "   'mar-may-2021/BKGDimage-localCues_Letter.png'          ,\n",
    "   'mar-may-2021/BKGDimage-localCues_LetterTex.png'       ,\n",
    "   'jun-jul-aug-nov-2021/BKGDimage_3LocalCues.png'        ,\n",
    "   'jun-jul-aug-nov-2021/BKGDimage_3LocalCues_Reverse.png',\n",
    "   'jun-jul-aug-nov-2021/BKGDimage_asymm.png'             ,\n",
    "   'jun-jul-aug-nov-2021/BKGDimage_asymmRev.png'          ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " *** processing 2022/BKGDimage-20220920.png\n",
      "   * rescaling and filtering\n",
      "   * finding center and shifting\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920_cropped.png\n",
      "   * finding hole and entrance positions\n",
      "     :::: NEW ARENA PICTURE INFO ::::\n",
      "       - cropped image extent (cm) = [-107.53483128834355, 86.71302147239263, -70.96838709677417, 75.12193548387097]\n",
      "       - arena center (cm) = [-3.126610429447865, 0.250645161290322]\n",
      "       - arena radius (cm) = 61.39457797\n",
      "      indices ::: (a) below -4 == targets;    (b) [-4,...,-1] == entrances;    (c) 1... == holes\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920_holes.txt\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920_holes.png\n",
      "   * cropping the BW version of the arena\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920_origin.png\n",
      " *** processing 2022/BKGDimage-20220920-probe2.png\n",
      "   * rescaling and filtering\n",
      "   * finding center and shifting\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920-probe2_cropped.png\n",
      "   * finding hole and entrance positions\n",
      "     :::: NEW ARENA PICTURE INFO ::::\n",
      "       - cropped image extent (cm) = [-102.67863496932516, 91.569217791411, -70.96838709677417, 75.12193548387097]\n",
      "       - arena center (cm) = [1.7295858895705294, 0.250645161290322]\n",
      "       - arena radius (cm) = 61.39457797\n",
      "      indices ::: (a) below -4 == targets;    (b) [-4,...,-1] == entrances;    (c) 1... == holes\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920-probe2_holes.txt\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920-probe2_holes.png\n",
      "   * cropping the BW version of the arena\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20220920-probe2_origin.png\n",
      " *** processing 2022/BKGDimage-20221104.png\n",
      "   * rescaling and filtering\n",
      "   * finding center and shifting\n",
      " ::: WARNING -> The arena is zoomed in/out! Trying to compensate that...\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20221104_cropped.png\n",
      "   * finding hole and entrance positions\n",
      "     :::: NEW ARENA PICTURE INFO ::::\n",
      "       - cropped image extent (cm) = [-105.53609805924413, 92.53703779366703, -72.46542805100178, 76.50943533697632]\n",
      "       - arena center (cm) = [0.9282124616956082, 0.1598178506375234]\n",
      "       - arena radius (cm) = 61.67543056\n",
      "      indices ::: (a) below -4 == targets;    (b) [-4,...,-1] == entrances;    (c) 1... == holes\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20221104_holes.txt\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20221104_holes.png\n",
      "   * cropping the BW version of the arena\n",
      "     ... saving: arena_picture\\2022/BKGDimage-20221104_origin.png\n"
     ]
    }
   ],
   "source": [
    "# calculating the center of the arena for the reference image\n",
    "# in pixels\n",
    "\n",
    "arena_ref_xyr    = edged.find_arena_in_image('',arena_ref_img_mark,arena_color=mark_scheme.arena_boundary)\n",
    "arena_center_ref = numpy.array(arena_ref_xyr[:2])\n",
    "#_,arena_center_ref = plib.get_center_radius_in_px(arena_pic_bounds.arena_pic_left,arena_pic_bounds.arena_pic_right,arena_pic_bounds.arena_pic_bottom,arena_pic_bounds.arena_pic_top,\\\n",
    "#                                         plib.get_arena_diameter_cm() / 2.0, plib.get_arena_center(), plib.get_arena_picture_file_width_height()[0], plib.get_arena_picture_file_width_height()[1])\n",
    "\n",
    "for do_processing,p,p_preproc,p_entrance,p_mark,dx,dy,contrast_factor,bright_factor,dx_correct,dy_correct,r_correct in arena_pic:\n",
    "    if not do_processing:\n",
    "        continue\n",
    "\n",
    "    has_picture_with_marks    = len(p_mark)>0\n",
    "    has_picture_with_entrance = len(p_entrance)>0\n",
    "    has_picture_preproc       = len(p_preproc)>0\n",
    "\n",
    "    if (not has_picture_with_marks) and not(arena_pic_name in mark_pic_not_required):\n",
    "        raise RuntimeError('ERROR! - You must make a mark picture for the arena you are trying to crop: %s' % arena_pic_name)\n",
    "\n",
    "    #p,dx,dy,contrast_factor,bright_factor = arena_pic[0]\n",
    "    print(' *** processing %s' % p)\n",
    "\n",
    "    arena_pic_name          = os.path.join(arena_pic_dir,p)\n",
    "    arena_pic_mark_name     = os.path.join(arena_pic_dir,p_mark)     if has_picture_with_marks    else ''\n",
    "    arena_pic_entrance_name = os.path.join(arena_pic_dir,p_entrance) if has_picture_with_entrance else ''\n",
    "    arena_pic_preproc_name  = os.path.join(arena_pic_dir,p_preproc)  if has_picture_preproc       else ''\n",
    "\n",
    "    cropped_pic_name    = os.path.join(arena_pic_dir,p.replace('.png','_cropped.png'))\n",
    "    holes_pic_filename  = os.path.join(arena_pic_dir,p.replace('.png','_holes.png'))\n",
    "    holes_txt_filename  = os.path.join(arena_pic_dir,p.replace('.png','_holes.txt'))\n",
    "    origin_pic_name     = os.path.join(arena_pic_dir,p.replace('.png','_origin.png'))\n",
    "    if delete_before_processing:\n",
    "        for f in [cropped_pic_name,holes_pic_filename,holes_txt_filename,origin_pic_name]:\n",
    "            if os.path.isfile(f):\n",
    "                os.remove(f)\n",
    "    \n",
    "    get_transformed_picture = lambda pic_name,radius_ref=1.0,radius=1.0: edged.transform_pic_to_match_distance(arena_pic_ref,pic_name,dx_ref,dy_ref,dx,dy,radius_ref,radius)\n",
    "    to_gray_enhanced        = lambda im: PIL.ImageEnhance.Contrast(PIL.ImageEnhance.Brightness(to_gray(im)).enhance(bright_factor)).enhance(contrast_factor)\n",
    "    erase_outside_arena     = lambda im,xyr: edged.erase_outside_circle(im, xyr[:2], xyr[2]+r_correct, bgcolor=1.0)\n",
    "    get_cropped_aligned_pic = lambda im,xyr: edged.align_arena_center_and_crop(arena_pic_ref,im,arena_center_ref,xyr[:2],bgcolor_rgb=None,dx_correct_px=dx_correct,dy_correct_px=dy_correct,xRange_world_ref=dx_ref,xRange_world=dx,yRange_world_ref=dy_ref,yRange_world=dy)\n",
    "    apply_all_transforms    = lambda pic_name,xyr,radius_ref=1.0,radius=1.0: get_cropped_aligned_pic(erase_outside_arena(to_gray_enhanced(get_transformed_picture(pic_name,radius_ref,radius)),xyr),xyr)[0]\n",
    "\n",
    "    # transforming the new arena pic to match the reference arena pic\n",
    "    print('   * rescaling and filtering')\n",
    "    img                  = get_transformed_picture(arena_pic_name) #edged.transform_pic_to_match_distance(arena_pic_ref,arena_pic_name,dx_ref,dy_ref,dx,dy)\n",
    "    img_high_contrast_bw = to_gray_enhanced(img)\n",
    "    img_mark             = get_transformed_picture(arena_pic_mark_name) if has_picture_with_marks else None\n",
    "\n",
    "    ################\n",
    "    ################\n",
    "    ################\n",
    "    # finding the center of the arena and shifting the new image to align the center with the ref image\n",
    "    ##\n",
    "    ##\n",
    "    print('   * finding center and shifting')\n",
    "    arena_xyr            = edged.find_arena_in_image(img_high_contrast_bw,img_mark,arena_color=mark_scheme.arena_boundary,dr=5,n_circle_points=100,circle_match_threshold=0.4)\n",
    "\n",
    "    # the arena picture being cropped is zoomed in/out\n",
    "    arena_radius_ref = 1.0\n",
    "    arena_radius     = 1.0\n",
    "    if abs(1.0 - arena_xyr[2]/arena_ref_xyr[2]) > arena_radius_tolerance:\n",
    "        print(' ::: WARNING -> The arena is zoomed in/out! Trying to compensate that...')\n",
    "        arena_radius_ref     = arena_ref_xyr[2]\n",
    "        arena_radius         = arena_xyr[2]\n",
    "        img                  = get_transformed_picture(arena_pic_name,arena_radius_ref,arena_radius) #edged.transform_pic_to_match_distance(arena_pic_ref,arena_pic_name,dx_ref,dy_ref,dx,dy)\n",
    "        img_high_contrast_bw = to_gray_enhanced(img)\n",
    "        img_mark             = get_transformed_picture(arena_pic_mark_name,arena_radius_ref,arena_radius) if has_picture_with_marks else None\n",
    "\n",
    "    # trying again with the proper transform\n",
    "    arena_xyr            = edged.find_arena_in_image(img_high_contrast_bw,mark_img=img_mark,arena_color=mark_scheme.arena_boundary,dr=5,n_circle_points=100,circle_match_threshold=0.4)\n",
    "    img_shift,img_extent = get_cropped_aligned_pic(img,arena_xyr)\n",
    "\n",
    "    # applying transforms to the supporting images\n",
    "    img_mark             = get_cropped_aligned_pic(img_mark,arena_xyr)[0] if has_picture_with_marks else None\n",
    "    img_entrance         = img_mark.copy() if has_picture_with_marks else (apply_all_transforms(arena_pic_entrance_name,arena_xyr,arena_radius_ref,arena_radius) if has_picture_with_entrance else None)\n",
    "    img_find_holes       = img_mark.copy() if has_picture_with_marks else (apply_all_transforms(arena_pic_preproc_name, arena_xyr,arena_radius_ref,arena_radius) if has_picture_preproc       else img_high_contrast_bw.copy())\n",
    "    \n",
    "    ################ saving\n",
    "    print('     ... saving: %s'%cropped_pic_name)\n",
    "    img_shift.save(cropped_pic_name,dpi=img.info['dpi'])\n",
    "    cropped_img_extent_str = '       - cropped image extent (cm) = ' + str(list(img_extent))\n",
    "\n",
    "    ################\n",
    "    ################\n",
    "    ################\n",
    "    # finding the holes of the arena\n",
    "    ##\n",
    "    ##\n",
    "    print('   * finding hole and entrance positions')\n",
    "\n",
    "    # setting up the transformation of coordinates to translate the holes and center into the world coordinates\n",
    "    T_cropped          = misc.LinearTransf2D( (0,img_shift.size[0]), img_extent[:2], (img_shift.size[1],0), img_extent[2:])\n",
    "    get_feat_xy_coord  = lambda xy_px: T_cropped(numpy.array(xy_px[:2]))\n",
    "    get_all_xy_coords  = lambda coord_px,k0: numpy.array([ (k+k0,*get_feat_xy_coord(h)) for k,h in enumerate(coord_px) ])\n",
    "\n",
    "    # finding entrances\n",
    "    entrance_px     = []                      # just dummy values to hold the place and prevent crash\n",
    "    entrance_coords = numpy.array([-1,-1,-1]) # just dummy values to hold the place and prevent crash\n",
    "    if has_picture_with_marks or has_picture_with_entrance:\n",
    "        entrance_color  = mark_scheme.entrances if has_picture_with_marks else None\n",
    "        entrance_px     = edged.find_clusters(img_entrance,threshold=0.5,find_less_than_threshold=True,color=entrance_color)\n",
    "        entrance_coords = get_all_xy_coords(entrance_px,-len(entrance_px))\n",
    "\n",
    "    # finding targets\n",
    "    targets_px     = []                      # just dummy values to hold the place and prevent crash\n",
    "    targets_coords = numpy.array([-1,-1,-1]) # just dummy values to hold the place and prevent crash\n",
    "    if has_picture_with_marks:\n",
    "        if mark_scheme.IsField('targets'):\n",
    "            target_color   = mark_scheme.targets if has_picture_with_marks else None\n",
    "            targets_px     = edged.find_clusters(img_mark,threshold=0.5,find_less_than_threshold=True,color=target_color)\n",
    "            targets_coords = get_all_xy_coords(targets_px,-len(entrance_px)-len(targets_px))\n",
    "        else:\n",
    "            print(' ***** to find the targets, you need a _mark.png picture from the original arena, marking the target with the mark_scheme.targets color')\n",
    "\n",
    "    # finding holes\n",
    "    holes_color  = mark_scheme.holes if has_picture_with_marks else None\n",
    "    holes_px     = edged.find_clusters(img_find_holes,threshold=0.5,find_less_than_threshold=True,color=holes_color)\n",
    "    holes_coords = get_all_xy_coords(holes_px,1)\n",
    "\n",
    "    # drawing holes in the picture\n",
    "    img_find_holes = img_find_holes.convert('RGB')\n",
    "    draw_result = PIL.ImageDraw.Draw(img_find_holes)\n",
    "    holes_color  = mark_scheme.holes if has_picture_with_marks else None\n",
    "    for k,h in enumerate(holes_px):\n",
    "        hole_radius = numpy.max((h[2],2))\n",
    "        draw_result.ellipse((h[0]-hole_radius,h[1]-hole_radius,h[0]+hole_radius,h[1]+hole_radius),outline=(255,0,0),fill=None)\n",
    "        draw_result.text(tuple(numpy.array(h[:2])+hole_radius/2),str(k+1),anchor='rb',fill=holes_color)\n",
    "    entrance_color  = mark_scheme.entrances if has_picture_with_marks else (0,255,0)\n",
    "    for k,h in enumerate(entrance_px):\n",
    "        hole_radius = numpy.max((h[2],2))\n",
    "        draw_result.ellipse((h[0]-hole_radius,h[1]-hole_radius,h[0]+hole_radius,h[1]+hole_radius),outline=entrance_color,fill=None)\n",
    "        draw_result.text(tuple(numpy.array(h[:2])+hole_radius/2),str(k-len(entrance_px)),anchor='rb',fill=entrance_color)\n",
    "    target_color   = mark_scheme.targets if has_picture_with_marks else (255,0,0)\n",
    "    for k,h in enumerate(targets_px):\n",
    "        hole_radius = numpy.max((h[2],2))\n",
    "        draw_result.ellipse((h[0]-hole_radius,h[1]-hole_radius,h[0]+hole_radius,h[1]+hole_radius),outline=target_color,fill=None)\n",
    "        draw_result.text(tuple(numpy.array(h[:2])+hole_radius/2),str(k-len(entrance_px)-len(targets_px)),anchor='rt',fill=target_color)\n",
    "    \n",
    "    # applying transformations to arena coords\n",
    "    arena_center_world = list(misc.LinearTransf2D( (0,img.size[0]), dx, (img.size[1],0), dy)(numpy.array(arena_xyr[:2])).flatten())\n",
    "    arena_R_world      = (misc.LinearTransf( (0.0,img.size[0]), (0.0,sum(numpy.abs(dx))) )(arena_xyr[2]) + misc.LinearTransf( (0.0,img.size[1]), (0.0,sum(numpy.abs(dy))) )(arena_xyr[2]))/2.0\n",
    "\n",
    "    # writing hole txt file and picture\n",
    "    print('     :::: NEW ARENA PICTURE INFO ::::')\n",
    "    arena_info_txt  = '%s\\n       - arena center (cm) = %s\\n       - arena radius (cm) = %.8f'%(cropped_img_extent_str,str(arena_center_world),arena_R_world)\n",
    "    arena_info_txt += '\\n      indices ::: (a) below {0} == targets;    (b) [{0},...,-1] == entrances;    (c) 1... == holes'.format(-len(entrance_px))\n",
    "    print(arena_info_txt)\n",
    "    header_str = 'Hole coordinates for file %s\\nCropped file = %s\\nHole visualizaton = %s\\n%s\\nhole index,x (cm),y (cm)'%(arena_pic_name,cropped_pic_name,holes_pic_filename,arena_info_txt)\n",
    "    print('     ... saving: %s'%holes_txt_filename)\n",
    "    numpy.savetxt(holes_txt_filename,numpy.row_stack((targets_coords,entrance_coords,holes_coords)),fmt=['%d','%.8f','%.8f'],delimiter=',',header=header_str)\n",
    "    img_find_holes.save(holes_pic_filename,dpi=img.info['dpi'])\n",
    "    print('     ... saving: %s'%holes_pic_filename)\n",
    "\n",
    "    # transforming the high contrast version of the picture (used to detect the arena) with the (0,0)_world in blue\n",
    "    print('   * cropping the BW version of the arena')\n",
    "    img_high_contrast_bw = img_high_contrast_bw.convert('RGB')\n",
    "    draw_result = PIL.ImageDraw.Draw(img_high_contrast_bw)\n",
    "    draw_result.ellipse((arena_xyr[0]-arena_xyr[2], arena_xyr[1]-arena_xyr[2], arena_xyr[0]+arena_xyr[2], arena_xyr[1]+arena_xyr[2]), outline=(255,0,0,0))\n",
    "    draw_result.ellipse((arena_xyr[0]-2,arena_xyr[1]-2,arena_xyr[0]+2,arena_xyr[1]+2),fill=(0,0,255))\n",
    "    img_high_contrast_bw_shift = get_cropped_aligned_pic(img_high_contrast_bw,arena_xyr)[0]\n",
    "    print('     ... saving: %s'%origin_pic_name)\n",
    "    img_high_contrast_bw_shift.save(origin_pic_name,dpi=img.info['dpi'])\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "570feb405e2e27c949193ac68f46852414290d515b0ba6e5d90d076ed2284471"
  },
  "kernelspec": {
   "display_name": "Python 3.8.2 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}