diff --git a/MindEyeV2/antspy/ants/core/__init__.py b/MindEyeV2/antspy/ants/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..665d08f77f6fdb4ea4b4d97c0b0a50bc52335c31 --- /dev/null +++ b/MindEyeV2/antspy/ants/core/__init__.py @@ -0,0 +1,42 @@ +from .ants_image_io import (image_header_info, + image_clone, + image_read, + dicom_read, + image_write, + make_image, + from_numpy, + from_numpy_like, + new_image_like) +from .ants_image import (ANTsImage, + copy_image_info, + set_origin, + get_origin, + set_direction, + get_direction, + set_spacing, + get_spacing, + is_image, + from_pointer) +from .ants_metric_io import (new_ants_metric, + create_ants_metric, + supported_metrics) +from .ants_transform_io import (create_ants_transform, + new_ants_transform, + read_transform, + write_transform, + transform_from_displacement_field, + transform_to_displacement_field, + fsl2antstransform) +from .ants_transform import (ANTsTransform, + set_ants_transform_parameters, + get_ants_transform_parameters, + get_ants_transform_fixed_parameters, + set_ants_transform_fixed_parameters, + apply_ants_transform, + apply_ants_transform_to_point, + apply_ants_transform_to_vector, + apply_ants_transform_to_image, + invert_ants_transform, + compose_ants_transforms, + transform_index_to_physical_point, + transform_physical_point_to_index) \ No newline at end of file diff --git a/MindEyeV2/antspy/ants/core/ants_image_io.py b/MindEyeV2/antspy/ants/core/ants_image_io.py new file mode 100644 index 0000000000000000000000000000000000000000..d6d349740d16949ea10c82a434586be5cdc7ddf2 --- /dev/null +++ b/MindEyeV2/antspy/ants/core/ants_image_io.py @@ -0,0 +1,529 @@ +""" +Image IO +""" + +__all__ = [ + "image_header_info", + "image_clone", + "image_read", + "dicom_read", + "image_write", + "make_image", + "from_numpy", + "from_numpy_like", + "new_image_like" +] + +import os +import json +import numpy as np +import warnings + +import ants +from ants.internal import get_lib_fn, short_ptype, infer_dtype +from ants.decorators import image_method + +_supported_pclasses = {"scalar", "vector", "rgb", "rgba","symmetric_second_rank_tensor"} +_supported_ptypes = {"unsigned char", "unsigned int", "float", "double"} +_supported_ntypes = {"uint8", "uint32", "float32", "float64"} +_unsupported_ptypes = {"char", "unsigned short", "short", "int"} +_unsupported_ptype_map = { + "char": "float", + "unsigned short": "unsigned int", + "short": "float", + "int": "float", +} +_image_type_map = {"scalar": "", "vector": "V", "rgb": "RGB", "rgba": "RGBA", "symmetric_second_rank_tensor": "SSRT" } +_ptype_type_map = { + "unsigned char": "UC", + "unsigned int": "UI", + "float": "F", + "double": "D", +} + +_ntype_type_map = {"uint8": "UC", "uint32": "UI", "float32": "F", "float64": "D"} +_npy_to_itk_map = { + "uint8": "unsigned char", + "uint32": "unsigned int", + "float32": "float", + "float64": "double", +} + +_image_read_dict = {} +for itype in {"scalar", "vector", "rgb", "rgba", "symmetric_second_rank_tensor"}: + _image_read_dict[itype] = {} + for p in _supported_ptypes: + _image_read_dict[itype][p] = {} + for d in {2, 3, 4}: + ita = _image_type_map[itype] + pa = _ptype_type_map[p] + _image_read_dict[itype][p][d] = "imageRead%s%s%i" % (ita, pa, d) + +def from_numpy( + data, origin=None, spacing=None, direction=None, has_components=False, is_rgb=False +): + """ + Create an ANTsImage object from a numpy array + + ANTsR function: `as.antsImage` + + Arguments + --------- + data : ndarray + image data array + + origin : tuple/list + image origin + + spacing : tuple/list + image spacing + + direction : list/ndarray + image direction + + has_components : boolean + whether the image has components + + Returns + ------- + ANTsImage + image with given data and any given information + """ + + # this is historic but should be removed once tests can pass without it + if data.dtype.name == 'float64': + data = data.astype('float32') + + # if dtype is not supported, cast to best available + best_dtype = infer_dtype(data.dtype) + if best_dtype != data.dtype: + data = data.astype(best_dtype) + + img = _from_numpy(data.T.copy(), origin, spacing, direction, has_components, is_rgb) + return img + + +def _from_numpy( + data, origin=None, spacing=None, direction=None, has_components=False, is_rgb=False +): + """ + Internal function for creating an ANTsImage + """ + if is_rgb: + has_components = True + ndim = data.ndim + if has_components: + ndim -= 1 + dtype = data.dtype.name + ptype = _npy_to_itk_map[dtype] + + data = np.array(data) + + if origin is None: + origin = tuple([0.0] * ndim) + if spacing is None: + spacing = tuple([1.0] * ndim) + if direction is None: + direction = np.eye(ndim) + + libfn = get_lib_fn("fromNumpy%s%i" % (_ntype_type_map[dtype], ndim)) + + if not has_components: + itk_image = libfn(data, data.shape[::-1]) + ants_image = ants.from_pointer(itk_image) + ants_image.set_origin(origin) + ants_image.set_spacing(spacing) + ants_image.set_direction(direction) + ants_image._ndarr = data + else: + arrays = [data[i, ...].copy() for i in range(data.shape[0])] + data_shape = arrays[0].shape + ants_images = [] + for i in range(len(arrays)): + tmp_ptr = libfn(arrays[i], data_shape[::-1]) + tmp_img = ants.from_pointer(tmp_ptr) + tmp_img.set_origin(origin) + tmp_img.set_spacing(spacing) + tmp_img.set_direction(direction) + tmp_img._ndarr = arrays[i] + ants_images.append(tmp_img) + ants_image = ants.merge_channels(ants_images) + if is_rgb: + ants_image = ants_image.vector_to_rgb() + return ants_image + + +def make_image( + imagesize, + voxval=0, + spacing=None, + origin=None, + direction=None, + has_components=False, + pixeltype="float", +): + """ + Make an image with given size and voxel value or given a mask and vector + + ANTsR function: `makeImage` + + Arguments + --------- + shape : tuple/ANTsImage + input image size or mask + + voxval : scalar + input image value or vector, size of mask + + spacing : tuple/list + image spatial resolution + + origin : tuple/list + image spatial origin + + direction : list/ndarray + direction matrix to convert from index to physical space + + components : boolean + whether there are components per pixel or not + + pixeltype : float + data type of image values + + Returns + ------- + ANTsImage + """ + if ants.is_image(imagesize): + img = imagesize.clone() + sel = imagesize > 0 + if voxval.ndim > 1: + voxval = voxval.flatten() + if (len(voxval) == int((sel > 0).sum())) or (len(voxval) == 0): + img[sel] = voxval + else: + raise ValueError( + "Num given voxels %i not same as num positive values %i in `imagesize`" + % (len(voxval), int((sel > 0).sum())) + ) + return img + else: + if isinstance(voxval, (tuple, list, np.ndarray)): + array = np.asarray(voxval).astype("float32").reshape(imagesize) + else: + array = np.full(imagesize, voxval, dtype="float32") + image = from_numpy( + array, + origin=origin, + spacing=spacing, + direction=direction, + has_components=has_components, + ) + return image.clone(pixeltype) + + +def image_header_info(filename): + """ + Read file info from image header + + ANTsR function: `antsImageHeaderInfo` + + Arguments + --------- + filename : string + name of image file from which info will be read + + Returns + ------- + dict + """ + if not os.path.exists(filename): + raise Exception("filename does not exist") + + libfn = get_lib_fn("antsImageHeaderInfo") + retval = libfn(filename) + retval["dimensions"] = tuple(retval["dimensions"]) + retval["origin"] = tuple([round(o, 4) for o in retval["origin"]]) + retval["spacing"] = tuple([round(s, 4) for s in retval["spacing"]]) + retval["direction"] = np.round(retval["direction"], 4) + return retval + +def image_clone(image, pixeltype=None): + """ + Clone an ANTsImage + + ANTsR function: `antsImageClone` + + Arguments + --------- + image : ANTsImage + image to clone + + dtype : string (optional) + new datatype for image + + Returns + ------- + ANTsImage + """ + return image.clone(pixeltype) + + +def image_read(filename, dimension=None, pixeltype="float", reorient=False): + """ + Read an ANTsImage from file + + ANTsR function: `antsImageRead` + + Arguments + --------- + filename : string + Name of the file to read the image from. + + dimension : int + Number of dimensions of the image read. This need not be the same as + the dimensions of the image in the file. Allowed values: 2, 3, 4. + If not provided, the dimension is obtained from the image file + + pixeltype : string + C++ datatype to be used to represent the pixels read. This datatype + need not be the same as the datatype used in the file. + Options: unsigned char, unsigned int, float, double + + reorient : boolean | string + if True, the image will be reoriented to RPI if it is 3D + if False, nothing will happen + if string, this should be the 3-letter orientation to which the + input image will reoriented if 3D. + if the image is 2D, this argument is ignored + + Returns + ------- + ANTsImage + """ + if filename.endswith(".npy"): + filename = os.path.expanduser(filename) + img_array = np.load(filename) + if os.path.exists(filename.replace(".npy", ".json")): + with open(filename.replace(".npy", ".json")) as json_data: + img_header = json.load(json_data) + ants_image = from_numpy( + img_array, + origin=img_header.get("origin", None), + spacing=img_header.get("spacing", None), + direction=np.asarray(img_header.get("direction", None)), + has_components=img_header.get("components", 1) > 1, + ) + else: + img_header = {} + ants_image = from_numpy(img_array) + + else: + filename = os.path.expanduser(filename) + if not os.path.exists(filename): + raise ValueError("File %s does not exist!" % filename) + + hinfo = image_header_info(filename) + ptype = hinfo["pixeltype"] + pclass = hinfo["pixelclass"] + ndim = hinfo["nDimensions"] + ncomp = hinfo["nComponents"] + is_rgb = False + if pclass == "rgb": + pclass = "vector" + if pclass == "rgba": + pclass = "vector" + if pclass == "symmetric_second_rank_tensor": + pclass = "vector" +# is_rgb = True if pclass == "rgb" else False + if dimension is not None: + ndim = dimension + + # error handling on pixelclass + if pclass not in _supported_pclasses: + raise ValueError("Pixel class %s not supported!" % pclass) + + # error handling on pixeltype + if ptype in _unsupported_ptypes: + ptype = _unsupported_ptype_map.get(ptype, "unsupported") + if ptype == "unsupported": + raise ValueError("Pixeltype %s not supported" % ptype) + + # error handling on dimension + if (ndim < 2) or (ndim > 4): + raise ValueError("Found %i-dimensional image - not supported!" % ndim) + + libfn = get_lib_fn(_image_read_dict[pclass][ptype][ndim]) + itk_pointer = libfn(filename) + + ants_image = ants.from_pointer(itk_pointer) + + if pixeltype is not None: + ants_image = ants_image.clone(pixeltype) + + if (reorient != False) and (ants_image.dimension == 3): + if reorient == True: + ants_image = ants_image.reorient_image2("RPI") + elif isinstance(reorient, str): + ants_image = ants_image.reorient_image2(reorient) + + return ants_image + + +def dicom_read(directory, pixeltype="float"): + """ + Read a set of dicom files in a directory into a single ANTsImage. + The origin of the resulting 3D image will be the origin of the + first dicom image read. + + Arguments + --------- + directory : string + folder in which all the dicom images exist + + Returns + ------- + ANTsImage + + Example + ------- + >>> import ants + >>> img = ants.dicom_read('~/desktop/dicom-subject/') + """ + slices = [] + imgidx = 0 + for imgpath in os.listdir(directory): + if imgpath.endswith(".dcm"): + if imgidx == 0: + tmp = image_read( + os.path.join(directory, imgpath), dimension=3, pixeltype=pixeltype + ) + origin = tmp.origin + spacing = tmp.spacing + direction = tmp.direction + tmp = tmp.numpy()[:, :, 0] + else: + tmp = image_read( + os.path.join(directory, imgpath), dimension=2, pixeltype=pixeltype + ).numpy() + + slices.append(tmp) + imgidx += 1 + + slices = np.stack(slices, axis=-1) + return from_numpy(slices, origin=origin, spacing=spacing, direction=direction) + +@image_method +def image_write(image, filename, ri=False): + """ + Write an ANTsImage to file + + ANTsR function: `antsImageWrite` + + Arguments + --------- + image : ANTsImage + image to save to file + + filename : string + name of file to which image will be saved + + ri : boolean + if True, return image. This allows for using this function in a pipeline: + >>> img2 = img.smooth_image(2.).image_write(file1, ri=True).threshold_image(0,20).image_write(file2, ri=True) + if False, do not return image + """ + if filename.endswith(".npy"): + img_array = image.numpy() + img_header = { + "origin": image.origin, + "spacing": image.spacing, + "direction": image.direction.tolist(), + "components": image.components, + } + + np.save(filename, img_array) + with open(filename.replace(".npy", ".json"), "w") as outfile: + json.dump(img_header, outfile) + else: + image.to_file(filename) + + if ri: + return image + +@image_method +def clone(image, pixeltype=None): + """ + Create a copy of the given ANTsImage with the same data and info, possibly with + a different data type for the image data. Only supports casting to + uint8 (unsigned char), uint32 (unsigned int), float32 (float), and float64 (double) + + Arguments + --------- + dtype: string (optional) + if None, the dtype will be the same as the cloned ANTsImage. Otherwise, + the data will be cast to this type. This can be a numpy type or an ITK + type. + Options: + 'unsigned char' or 'uint8', + 'unsigned int' or 'uint32', + 'float' or 'float32', + 'double' or 'float64' + + Returns + ------- + ANTsImage + """ + if pixeltype is None: + pixeltype = image.pixeltype + + if pixeltype not in _supported_ptypes: + raise ValueError('Pixeltype %s not supported. Supported types are %s' % (pixeltype, _supported_ptypes)) + + if image.has_components and (not image.is_rgb): + comp_imgs = ants.split_channels(image) + comp_imgs_cloned = [comp_img.clone(pixeltype) for comp_img in comp_imgs] + return ants.merge_channels(comp_imgs_cloned, channels_first=image.channels_first) + else: + p1_short = short_ptype(image.pixeltype) + p2_short = short_ptype(pixeltype) + ndim = image.dimension + fn_suffix = '%s%i' % (p2_short,ndim) + libfn = get_lib_fn('antsImageClone%s'%fn_suffix) + pointer_cloned = libfn(image.pointer) + return ants.from_pointer(pointer_cloned) + +copy = clone + +@image_method +def new_image_like(image, data): + """ + Create a new ANTsImage with the same header information, but with + a new image array. + + Arguments + --------- + data : ndarray or py::capsule + New array or pointer for the image. + It must have the same shape as the current + image data. + + Returns + ------- + ANTsImage + """ + if not isinstance(data, np.ndarray): + raise ValueError('data must be a numpy array') + if not image.has_components: + if data.shape != image.shape: + raise ValueError('given array shape (%s) and image array shape (%s) do not match' % (data.shape, image.shape)) + else: + if (data.shape[-1] != image.components) or (data.shape[:-1] != image.shape): + raise ValueError('given array shape (%s) and image array shape (%s) do not match' % (data.shape[1:], image.shape)) + + return from_numpy(data, origin=image.origin, + spacing=image.spacing, direction=image.direction, + has_components=image.has_components) + +def from_numpy_like(data, image): + return new_image_like(image, data) \ No newline at end of file diff --git a/MindEyeV2/antspy/ants/math/__init__.py b/MindEyeV2/antspy/ants/math/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9f920dcd2213ca95db97dd4da79cc17ffe112265 --- /dev/null +++ b/MindEyeV2/antspy/ants/math/__init__.py @@ -0,0 +1,14 @@ +from .averaging import average_images +from .get_centroids import get_centroids +from .get_neighborhood import get_neighborhood_in_mask, get_neighborhood_at_voxel +from .hausdorff_distance import hausdorff_distance +from .image_similarity import image_similarity +from .metrics import image_mutual_information +from .quantile import (ilr, + rank_intensity, + quantile, + regress_poly, + regress_components, + get_average_of_timeseries, + compcor, + bandpass_filter_matrix) \ No newline at end of file diff --git a/MindEyeV2/antspy/ants/math/averaging.py b/MindEyeV2/antspy/ants/math/averaging.py new file mode 100644 index 0000000000000000000000000000000000000000..b7b369719ab03c8ba6a80fc2da07d9b22d2b2e39 --- /dev/null +++ b/MindEyeV2/antspy/ants/math/averaging.py @@ -0,0 +1,100 @@ +import os +from tempfile import mktemp + +import numpy as np + +import ants + +__all__ = ['average_images'] + + +def average_images( x, normalize=True, mask=None, imagetype=0, sum_image_threshold=3, return_sum_image=False, verbose=False ): + """ + average a list of images + + images will be resampled automatically to the largest image space; + this is not a registration so images should be in the same physical + space to begin with. + + x : a list containing either filenames or antsImages + + normalize : boolean + + mask : None or integer; this will perform a masked averaging which can + be useful when images have only partial coverage. integer greater + than zero will perform morphological closing. + + imagetype : integer + choose 0/1/2/3 mapping to scalar/vector/tensor/time-series + + sum_image_threshold : integer + only average regions with overlap greater than or equal to this value + + return_sum_image : boolean + returns the average and the image that show ROI overlap; primarily for debugging + + verbose : boolean + will print progress + + Returns + ------- + ANTsImage + + Example + ------- + >>> import ants + >>> x0=[ ants.get_data('r16'), ants.get_data('r27'), ants.get_data('r62'), ants.get_data('r64') ] + >>> x1=[] + >>> for k in range(len(x0)): + >>> x1.append( ants.image_read( x0[k] ) ) + >>> avg=ants.average_images(x0) + >>> avg1=ants.average_images(x1) + >>> avg2=ants.average_images(x1,mask=0) + >>> avg3=ants.average_images(x1,mask=1,normalize=True) + """ + import numpy as np + + def gli( y, normalize=False ): + if isinstance(y,str): + y=ants.image_read(y) + if normalize: + y=y/y.mean() + return y + + biggest=0 + biggestind=0 + for k in range( len( x ) ): + locimg = gli( x[k], False ) + sz=np.prod( locimg.shape ) + if sz > biggest: + biggest=sz + biggestind=k + + avg = gli( x[biggestind], False ) * 0 + scl = float( 1.0 / len(x)) + if mask is not None: + sumimg = gli( x[biggestind], False ) * 0 + + for k in range( len( x ) ): + if verbose and k % 20 == 0: + print( str(k)+'...', end='',flush=True) + locimg = gli( x[k], normalize ) + temp = ants.resample_image_to_target( locimg, avg, interp_type='linear', imagetype=imagetype ) + avg = avg + temp + if mask is not None: + fgmask = ants.threshold_image(temp,'Otsu',1) + if mask > 0: + fgmask = ants.morphology(fgmask,"close",mask) + sumimg = sumimg + fgmask + + if return_sum_image: + return avg * scl, sumimg + if mask is None: + avg = avg * scl + else: + nonzero = sumimg > sum_image_threshold + tozero = sumimg <= sum_image_threshold + avg[nonzero] = avg[nonzero] / sumimg[nonzero] + avg[tozero] = 0 + return avg + diff --git a/MindEyeV2/antspy/ants/math/get_centroids.py b/MindEyeV2/antspy/ants/math/get_centroids.py new file mode 100644 index 0000000000000000000000000000000000000000..0aa8020bf6dcdc63f44e6113722bac9f62849fb7 --- /dev/null +++ b/MindEyeV2/antspy/ants/math/get_centroids.py @@ -0,0 +1,58 @@ +__all__ = ["get_centroids"] + +import numpy as np +import ants +from ants.decorators import image_method + +@image_method +def get_centroids(image, clustparam=0): + """ + Reduces a variate/statistical/network image to a set of centroids + describing the center of each stand-alone non-zero component in the image + + ANTsR function: `getCentroids` + + Arguments + --------- + image : ANTsImage + image from which centroids will be calculated + + clustparam : integer + look at regions greater than or equal to this size + + Returns + ------- + ndarray + + Example + ------- + >>> import ants + >>> image = ants.image_read( ants.get_ants_data( "r16" ) ) + >>> image = ants.threshold_image( image, 90, 120 ) + >>> image = ants.label_clusters( image, 10 ) + >>> cents = ants.get_centroids( image ) + """ + imagedim = image.dimension + if clustparam > 0: + mypoints = ants.label_clusters(image, clustparam, max_thresh=1e15) + if clustparam == 0: + mypoints = image.clone() + mypoints = ants.label_stats(mypoints, mypoints) + nonzero = mypoints[["LabelValue"]] > 0 + mypoints = mypoints[nonzero["LabelValue"]] + mypoints = mypoints.iloc[:, :] + x = mypoints.x + y = mypoints.y + + if imagedim == 3: + z = mypoints.z + else: + z = np.zeros(mypoints.shape[0]) + + if imagedim == 4: + t = mypoints.t + else: + t = np.zeros(mypoints.shape[0]) + + centroids = np.stack([x, y, z, t]).T + return centroids diff --git a/MindEyeV2/antspy/ants/math/get_neighborhood.py b/MindEyeV2/antspy/ants/math/get_neighborhood.py new file mode 100644 index 0000000000000000000000000000000000000000..0c1cb6c817f2770d73560b29b0eb18b14328eba6 --- /dev/null +++ b/MindEyeV2/antspy/ants/math/get_neighborhood.py @@ -0,0 +1,187 @@ + +__all__ = ['get_neighborhood_in_mask', + 'get_neighborhood_at_voxel'] + +import numpy as np + +import ants + +from ants.internal import get_lib_fn +from ants.decorators import image_method + +@image_method +def get_neighborhood_in_mask(image, mask, radius, physical_coordinates=False, + boundary_condition=None, spatial_info=False, get_gradient=False): + """ + Get neighborhoods for voxels within mask. + + This converts a scalar image to a matrix with rows that contain neighbors + around a center voxel + + ANTsR function: `getNeighborhoodInMask` + + Arguments + --------- + image : ANTsImage + image to get values from + + mask : ANTsImage + image indicating which voxels to examine. Each voxel > 0 will be used as the + center of a neighborhood + + radius : tuple/list + array of values for neighborhood radius (in voxels) + + physical_coordinates : boolean + whether voxel indices and offsets should be in voxel or physical coordinates + + boundary_condition : string (optional) + how to handle voxels in a neighborhood, but not in the mask. + None : fill values with `NaN` + `image` : use image value, even if not in mask + `mean` : use mean of all non-NaN values for that neighborhood + + spatial_info : boolean + whether voxel locations and neighborhood offsets should be returned along with pixel values. + + get_gradient : boolean + whether a matrix of gradients (at the center voxel) should be returned in + addition to the value matrix (WIP) + + Returns + ------- + if spatial_info is False: + if get_gradient is False: + ndarray + an array of pixel values where the number of rows is the size of the + neighborhood and there is a column for each voxel + + else if get_gradient is True: + dictionary w/ following key-value pairs: + values : ndarray + array of pixel values where the number of rows is the size of the + neighborhood and there is a column for each voxel. + + gradients : ndarray + array providing the gradients at the center voxel of each + neighborhood + + else if spatial_info is True: + dictionary w/ following key-value pairs: + values : ndarray + array of pixel values where the number of rows is the size of the + neighborhood and there is a column for each voxel. + + indices : ndarray + array provinding the center coordinates for each neighborhood + + offsets : ndarray + array providing the offsets from center for each voxel in a neighborhood + + Example + ------- + >>> import ants + >>> r16 = ants.image_read(ants.get_ants_data('r16')) + >>> mask = ants.get_mask(r16) + >>> mat = ants.get_neighborhood_in_mask(r16, mask, radius=(2,2)) + """ + if not ants.is_image(image): + raise ValueError('image must be ANTsImage type') + if not ants.is_image(mask): + raise ValueError('mask must be ANTsImage type') + if isinstance(radius, (int, float)): + radius = [radius]*image.dimension + if (not isinstance(radius, (tuple,list))) or (len(radius) != image.dimension): + raise ValueError('radius must be tuple or list with length == image.dimension') + + boundary = 0 + if boundary_condition == 'image': + boundary = 1 + elif boundary_condition == 'mean': + boundary = 2 + + libfn = get_lib_fn('getNeighborhoodMatrix%s' % image._libsuffix) + retvals = libfn(image.pointer, + mask.pointer, + list(radius), + int(physical_coordinates), + int(boundary), + int(spatial_info), + int(get_gradient)) + + if not spatial_info: + if get_gradient: + retvals['values'] = np.asarray(retvals['values']) + retvals['gradients'] = np.asarray(retvals['gradients']) + else: + retvals = np.asarray(retvals['matrix']) + else: + retvals['values'] = np.asarray(retvals['values']) + retvals['indices'] = np.asarray(retvals['indices']) + retvals['offsets'] = np.asarray(retvals['offsets']) + + return retvals + +@image_method +def get_neighborhood_at_voxel(image, center, kernel, physical_coordinates=False): + """ + Get a hypercube neighborhood at a voxel. Get the values in a local + neighborhood of an image. + + ANTsR function: `getNeighborhoodAtVoxel` + + Arguments + --------- + image : ANTsImage + image to get values from. + + center : tuple/list + indices for neighborhood center + + kernel : tuple/list + either a collection of values for neighborhood radius (in voxels) or + a binary collection of the same dimension as the image, specifying the shape of the neighborhood to extract + + physical_coordinates : boolean + whether voxel indices and offsets should be in voxel + or physical coordinates + + Returns + ------- + dictionary w/ following key-value pairs: + values : ndarray + array of neighborhood values at the voxel + + indices : ndarray + matrix providing the coordinates for each value + + Example + ------- + >>> import ants + >>> img = ants.image_read(ants.get_ants_data('r16')) + >>> center = (2,2) + >>> radius = (3,3) + >>> retval = ants.get_neighborhood_at_voxel(img, center, radius) + """ + if not ants.is_image(image): + raise ValueError('image must be ANTsImage type') + + if (not isinstance(center, (tuple,list))) or (len(center) != image.dimension): + raise ValueError('center must be tuple or list with length == image.dimension') + + if (not isinstance(kernel, (tuple,list))) or (len(kernel) != image.dimension): + raise ValueError('kernel must be tuple or list with length == image.dimension') + + radius = [int((k-1)/2) for k in kernel] + + libfn = get_lib_fn('getNeighborhood%s' % image._libsuffix) + retvals = libfn(image.pointer, + list(center), + list(kernel), + list(radius), + int(physical_coordinates)) + for k in retvals.keys(): + retvals[k] = np.asarray(retvals[k]) + return retvals + + diff --git a/MindEyeV2/antspy/ants/math/hausdorff_distance.py b/MindEyeV2/antspy/ants/math/hausdorff_distance.py new file mode 100644 index 0000000000000000000000000000000000000000..1e604ff048c79a379781bf28329efa94ba85a712 --- /dev/null +++ b/MindEyeV2/antspy/ants/math/hausdorff_distance.py @@ -0,0 +1,40 @@ +__all__ = ["hausdorff_distance"] + +from ants.decorators import image_method +from ants.internal import get_lib_fn + +@image_method +def hausdorff_distance(image1, image2): + """ + Get Hausdorff distance between non-zero pixels in two images + + ANTsR function: `hausdorffDistance` + + Arguments + --------- + source image : ANTsImage + Source image + + target_image : ANTsImage + Target image + + Returns + ------- + data frame with "Distance" and "AverageDistance" + + Example + ------- + >>> import ants + >>> r16 = ants.image_read( ants.get_ants_data('r16') ) + >>> r64 = ants.image_read( ants.get_ants_data('r64') ) + >>> s16 = ants.kmeans_segmentation( r16, 3 )['segmentation'] + >>> s64 = ants.kmeans_segmentation( r64, 3 )['segmentation'] + >>> stats = ants.hausdorff_distance(s16, s64) + """ + image1_int = image1.clone("unsigned int") + image2_int = image2.clone("unsigned int") + + libfn = get_lib_fn("hausdorffDistance%iD" % image1_int.dimension) + d = libfn(image1_int.pointer, image2_int.pointer) + + return d diff --git a/MindEyeV2/antspy/ants/math/image_similarity.py b/MindEyeV2/antspy/ants/math/image_similarity.py new file mode 100644 index 0000000000000000000000000000000000000000..21edfc5451716acbc6016b9639122045f9416b37 --- /dev/null +++ b/MindEyeV2/antspy/ants/math/image_similarity.py @@ -0,0 +1,67 @@ + + +__all__ = ['image_similarity'] + +import ants +from ants.decorators import image_method + +@image_method +def image_similarity(fixed_image, moving_image, metric_type='MeanSquares', + fixed_mask=None, moving_mask=None, + sampling_strategy='regular', sampling_percentage=1.): + """ + Measure similarity between two images. + NOTE: Similarity is actually returned as distance (i.e. dissimilarity) + per ITK/ANTs convention. E.g. using Correlation metric, the similarity + of an image with itself returns -1. + + ANTsR function: `imageSimilarity` + + Arguments + --------- + fixed : ANTsImage + the fixed image + + moving : ANTsImage + the moving image + + metric_type : string + image metric to calculate + MeanSquares + Correlation + ANTSNeighborhoodCorrelation + MattesMutualInformation + JointHistogramMutualInformation + Demons + + fixed_mask : ANTsImage (optional) + mask for the fixed image + + moving_mask : ANTsImage (optional) + mask for the moving image + + sampling_strategy : string (optional) + sampling strategy, default is full sampling + None (Full sampling) + random + regular + + sampling_percentage : scalar + percentage of data to sample when calculating metric + Must be between 0 and 1 + + Returns + ------- + scalar + + Example + ------- + >>> import ants + >>> x = ants.image_read(ants.get_ants_data('r16')) + >>> y = ants.image_read(ants.get_ants_data('r30')) + >>> metric = ants.image_similarity(x,y,metric_type='MeanSquares') + """ + metric = ants.create_ants_metric(fixed_image, moving_image, metric_type, fixed_mask, + moving_mask, sampling_strategy, sampling_percentage) + return metric.get_value() + diff --git a/MindEyeV2/antspy/ants/math/metrics.py b/MindEyeV2/antspy/ants/math/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..fbeb9a03f8f98a8ab3dc0c502e358afe5966e316 --- /dev/null +++ b/MindEyeV2/antspy/ants/math/metrics.py @@ -0,0 +1,43 @@ + + + +__all__ = ['image_mutual_information'] + + +from ants.decorators import image_method +from ants.internal import get_lib_fn + +@image_method +def image_mutual_information(image1, image2): + """ + Compute mutual information between two ANTsImage types + + ANTsR function: `antsImageMutualInformation` + + Arguments + --------- + image1 : ANTsImage + image 1 + + image2 : ANTsImage + image 2 + + Returns + ------- + scalar + + Example + ------- + >>> import ants + >>> fi = ants.image_read( ants.get_ants_data('r16') ).clone('float') + >>> mi = ants.image_read( ants.get_ants_data('r64') ).clone('float') + >>> mival = ants.image_mutual_information(fi, mi) # -0.1796141 + """ + if (image1.pixeltype != 'float') or (image2.pixeltype != 'float'): + raise ValueError('Both images must have float pixeltype') + + if image1.dimension != image2.dimension: + raise ValueError('Both images must have same dimension') + + libfn = get_lib_fn('antsImageMutualInformation%iD' % image1.dimension) + return libfn(image1.pointer, image2.pointer) diff --git a/MindEyeV2/antspy/ants/math/quantile.py b/MindEyeV2/antspy/ants/math/quantile.py new file mode 100644 index 0000000000000000000000000000000000000000..daca1b48d5a01f5dc0ec3e832139a164718484bb --- /dev/null +++ b/MindEyeV2/antspy/ants/math/quantile.py @@ -0,0 +1,447 @@ + +__all__ = ['ilr', + 'rank_intensity', + 'quantile', + 'regress_poly', + 'regress_components', + 'get_average_of_timeseries', + 'compcor', + 'bandpass_filter_matrix' ] + +import numpy as np +from numpy.polynomial import Legendre +from scipy import linalg +from scipy.stats import pearsonr +from scipy.stats import rankdata +import pandas as pd +from pandas import DataFrame +import statsmodels.api as sm +import statsmodels.formula.api as smf + +import ants +from ants.decorators import image_method + +def rank_intensity( x, mask=None, get_mask=True, method='max', ): + """ + Rank transform the intensity of the input image with or without masking. + Intensities will transform from [0,1,2,55] to [0,1,2,3] so this may not be + appropriate for quantitative images - however, you never know. rank + transformations generally improve robustness so it is an empirical question + that should be evaluated. + + Arguments + --------- + + x : ANTsImage + input image + + mask : ANTsImage + optional mask + + get_mask: boolean + will estimate a mask when none provided + + method : a scipy rank method (max,min,average,dense) + + + return: transformed image + + Example + ------- + >>> img = ants.image_read(ants.get_data('r16')) + >>> ants.rank_intensity(img) + """ + if mask is not None: + fir = rankdata( (x*mask).numpy(), method=method ) + elif mask is None and get_mask == True: + mask = ants.get_mask( x ) + fir = rankdata( (x*mask).numpy(), method=method ) + else: + fir = rankdata( x.numpy(), method=method ) + fir = fir - 1 + fir = fir.reshape( x.shape ) + rimg = ants.from_numpy( fir.astype(float) ) + rimg = ants.iMath(rimg,"Normalize") + ants.copy_image_info( x, rimg ) + if mask is not None: + rimg = rimg * mask + return( rimg ) + + +def ilr( data_frame, voxmats, ilr_formula, verbose = False ): + """ + Image-based linear regression. + + This function simplifies calculating p-values from linear models + in which there is a similar formula that is applied many times + with a change in image-based predictors. Image-based variables + are stored in the input matrix list. They should be named + consistently in the input formula and in the image list. If they + are not, an error will be thrown. All input matrices should have + the same number of rows and columns. + + This function takes advantage of statsmodels R-style formulas. + + ANTsR function: `ilr` + + Arguments + --------- + + data_frame: This data frame contains all relevant predictors except for + the matrices associated with the image variables. One should convert + any categorical predictors ahead of time using `pd.get_dummies`. + + voxmats: The named list of matrices that contains the changing + predictors. + + ilr_formula: This is a character string that defines a valid regression + formula in the R-style. + + verbose: will print a little bit of diagnostic information that allows + a degree of model checking + + Returns + ------- + + A list of different matrices that contain names derived from the + formula and the coefficients of the regression model. The size of + the output values ( p-values, t-values, parameter values ) will match + the input matrix and, as such, can be converted to an image via `make_image` + + Example + ------- + + >>> nsub = 20 + >>> mu, sigma = 0, 1 + >>> outcome = np.random.normal( mu, sigma, nsub ) + >>> covar = np.random.normal( mu, sigma, nsub ) + >>> mat = np.random.normal( mu, sigma, (nsub, 500 ) ) + >>> mat2 = np.random.normal( mu, sigma, (nsub, 500 ) ) + >>> data = {'covar':covar,'outcome':outcome} + >>> df = pd.DataFrame( data ) + >>> vlist = { "mat1": mat, "mat2": mat2 } + >>> myform = " outcome ~ covar * mat1 " + >>> result = ants.ilr( df, vlist, myform) + >>> myform = " mat2 ~ covar + mat1 " + >>> result = ants.ilr( df, vlist, myform) + + """ + + nvoxmats = len( voxmats ) + if nvoxmats < 1 : + raise ValueError('Pass at least one matrix to voxmats list') + keylist = list(voxmats.keys()) + firstmat = keylist[0] + voxshape = voxmats[firstmat].shape + nvox = voxshape[1] + nmats = len( keylist ) + for k in keylist: + if voxmats[firstmat].shape != voxmats[k].shape: + raise ValueError('Matrices must have same number of rows (samples)') + + # test voxel + vox = 0 + nrows = data_frame.shape[0] + data_frame_vox = data_frame.copy() + for k in range( nmats ): + data = {keylist[k]: np.random.normal(0,1,nrows) } + temp = pd.DataFrame( data ) + data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 ) + mod = smf.ols(formula=ilr_formula, data=data_frame_vox ) + res = mod.fit() + modelNames = res.model.exog_names + if verbose: + print( data_frame_vox ) + print(res.summary()) + nOutcomes = len( modelNames ) + tValsOut = list() + pValsOut = list() + bValsOut = list() + for k in range( len( modelNames ) ): + bValsOut.append( np.zeros( nvox ) ) + pValsOut.append( np.zeros( nvox ) ) + tValsOut.append( np.zeros( nvox ) ) + + data_frame_vox = data_frame.copy() + for v in range( nmats ): + data = {keylist[v]: voxmats[keylist[v]][:,k] } + temp = pd.DataFrame( data ) + data_frame_vox = pd.concat([data_frame_vox.reset_index(drop=True),temp], axis=1 ) + for k in range( nvox ): + # first get the correct data frame + for v in range( nmats ): + data_frame_vox[ keylist[v] ] = voxmats[keylist[v]][:,k] + # then get the local model results + mod = smf.ols(formula=ilr_formula, data=data_frame_vox ) + res = mod.fit() + tvals = res.tvalues + pvals = res.pvalues + bvals = res.params + for v in range( len( modelNames ) ): + bValsOut[v][k] = bvals.iloc[v] + pValsOut[v][k] = pvals.iloc[v] + tValsOut[v][k] = tvals.iloc[v] + + bValsOutDict = { } + tValsOutDict = { } + pValsOutDict = { } + for v in range( len( modelNames ) ): + bValsOutDict[ 'coef_' + modelNames[v] ] = bValsOut[v] + tValsOutDict[ 'tval_' + modelNames[v] ] = tValsOut[v] + pValsOutDict[ 'pval_' + modelNames[v] ] = pValsOut[v] + + return { + 'modelNames': modelNames, + 'coefficientValues': bValsOutDict, + 'pValues': pValsOutDict, + 'tValues': tValsOutDict } + + +@image_method +def quantile(image, q, nonzero=True): + """ + Get the quantile values from an ANTsImage + + Examples + -------- + >>> img = ants.image_read(ants.get_data('r16')) + >>> ants.quantile(img, 0.5) + >>> ants.quantile(img, (0.5, 0.75)) + """ + img_arr = image.numpy() + if isinstance(q, (list,tuple)): + q = [qq*100. if qq <= 1. else qq for qq in q] + if nonzero: + img_arr = img_arr[img_arr>0] + vals = [np.percentile(img_arr, qq) for qq in q] + return tuple(vals) + elif isinstance(q, (float,int)): + if q <= 1.: + q = q*100. + if nonzero: + img_arr = img_arr[img_arr>0] + return np.percentile(img_arr[img_arr>0], q) + else: + raise ValueError('q argument must be list/tuple or float/int') + + +def regress_poly(degree, data, remove_mean=True, axis=-1): + """ + Returns data with degree polynomial regressed out. + :param bool remove_mean: whether or not demean data (i.e. degree 0), + :param int axis: numpy array axes along which regression is performed + """ + timepoints = data.shape[0] + # Generate design matrix + X = np.ones((timepoints, 1)) # quick way to calc degree 0 + for i in range(degree): + polynomial_func = Legendre.basis(i + 1) + value_array = np.linspace(-1, 1, timepoints) + X = np.hstack((X, polynomial_func(value_array)[:, np.newaxis])) + non_constant_regressors = X[:, :-1] if X.shape[1] > 1 else np.array([]) + betas = np.linalg.pinv(X).dot(data) + if remove_mean: + datahat = X.dot(betas) + else: # disregard the first layer of X, which is degree 0 + datahat = X[:, 1:].dot(betas[1:, ...]) + regressed_data = data - datahat + return regressed_data, non_constant_regressors + +def regress_components( data, components, remove_mean=True ): + """ + Returns data with components regressed out. + :param bool remove_mean: whether or not demean data (i.e. degree 0), + :param int axis: numpy array axes along which regression is performed + """ + timepoints = data.shape[0] + betas = np.linalg.pinv(components).dot(data) + if remove_mean: + datahat = components.dot(betas) + else: # disregard the first layer of X, which is degree 0 + datahat = components[:, 1:].dot(betas[1:, ...]) + regressed_data = data - datahat + return regressed_data + + +def get_average_of_timeseries( image, idx=None ): + """Average the timeseries into a dimension-1 image. + image: input time series image + idx: indices over which to average + """ + imagedim = image.dimension + if idx is None: + idx = range( image.shape[ imagedim - 1 ] ) + i0 = ants.slice_image( image, axis=image.dimension-1, idx=idx[0] ) * 0 + wt = 1.0 / len( idx ) + for k in idx: + i0 = i0 + ants.slice_image( image, axis=image.dimension-1, idx=k ) * wt + return( i0 ) + +def bandpass_filter_matrix( matrix, + tr=1, lowf=0.01, highf=0.1, order = 3): + """ + Bandpass filter the input time series image + + ANTsR function: `frequencyFilterfMRI` + + Arguments + --------- + + image: input time series image + + tr: sampling time interval (inverse of sampling rate) + + lowf: low frequency cutoff + + highf: high frequency cutoff + + order: order of the butterworth filter run using `filtfilt` + + Returns + ------- + filtered matrix + + Example + ------- + + >>> import numpy as np + >>> import ants + >>> import matplotlib.pyplot as plt + >>> brainSignal = np.random.randn( 400, 1000 ) + >>> tr = 1 + >>> filtered = ants.bandpass_filter_matrix( brainSignal, tr = tr ) + >>> nsamples = brainSignal.shape[0] + >>> t = np.linspace(0, tr*nsamples, nsamples, endpoint=False) + >>> k = 20 + >>> plt.plot(t, brainSignal[:,k], label='Noisy signal') + >>> plt.plot(t, filtered[:,k], label='Filtered signal') + >>> plt.xlabel('time (seconds)') + >>> plt.grid(True) + >>> plt.axis('tight') + >>> plt.legend(loc='upper left') + >>> plt.show() + """ + from scipy.signal import butter, filtfilt + + def butter_bandpass(lowcut, highcut, fs, order ): + nyq = 0.5 * fs + low = lowcut / nyq + high = highcut / nyq + b, a = butter(order, [low, high], btype='band') + return b, a + + def butter_bandpass_filter(data, lowcut, highcut, fs, order ): + b, a = butter_bandpass(lowcut, highcut, fs, order=order) + y = filtfilt(b, a, data) + return y + + fs = 1/tr # sampling rate based on tr + nsamples = matrix.shape[0] + ncolumns = matrix.shape[1] + matrixOut = matrix.copy() + for k in range( ncolumns ): + matrixOut[:,k] = butter_bandpass_filter( + matrix[:,k], lowf, highf, fs, order=order ) + return matrixOut + +def clean_data(arr, standardize=True): + """ + Remove columns from a NumPy array that have no variation or contain NA/Inf values. + Optionally standardize the remaining data. + + :param arr: NumPy array to be cleaned. + :param standardize: Boolean, if True standardize the data. + :return: Cleaned (and optionally standardized) NumPy array. + """ + valid_columns = [] + + for i in range(arr.shape[1]): + column = arr[:, i] + if np.any(column != column[0]) and not np.any(np.isnan(column)) and not np.any(np.isinf(column)): + valid_columns.append(i) + + cleaned_data = arr[:, valid_columns] + + if standardize: + mean = np.mean(cleaned_data, axis=0) + std_dev = np.std(cleaned_data, axis=0) + # Avoid division by zero in case of zero standard deviation + std_dev[std_dev == 0] = 1 + cleaned_data = (cleaned_data - mean) / std_dev + + return cleaned_data + +def compcor( boldImage, ncompcor=4, quantile=0.975, mask=None, filter_type=False, degree=2 ): + """ + Compute noise components from the input image + + ANTsR function: `compcor` + + this is adapted from nipy code https://github.com/nipy/nipype/blob/e29ac95fc0fc00fedbcaa0adaf29d5878408ca7c/nipype/algorithms/confounds.py + + Arguments + --------- + + boldImage: input time series image + + ncompcor: number of noise components to return + + quantile: quantile defining high-variance + + mask: mask defining brain or specific tissues + + filter_type: type off filter to apply to time series before computing + noise components. + + 'polynomial' - Legendre polynomial basis + False - None (mean-removal only) + + degree: order of polynomial used to remove trends from the timeseries + + Returns + ------- + dictionary containing: + + components: a numpy array + + basis: a numpy array containing the (non-constant) filter regressors + + Example + ------- + >>> cc = ants.compcor( ants.image_read(ants.get_ants_data("ch2")) ) + + """ + + def compute_tSTD(M, quantile, x=0, axis=0): + stdM = np.std(M, axis=axis) + # set bad values to x + stdM[stdM == 0] = x + stdM[np.isnan(stdM)] = x + tt = round( quantile*100 ) + threshold_std = np.percentile( stdM, tt ) + # threshold_std = quantile( stdM, quantile ) + return { 'tSTD': stdM, 'threshold_std': threshold_std} + if mask is None: + temp = ants.slice_image( boldImage, axis=boldImage.dimension-1, idx=0 ) + mask = ants.get_mask( temp ) + imagematrix = ants.timeseries_to_matrix( boldImage, mask ) + temp = compute_tSTD( imagematrix, quantile, 0 ) + tsnrmask = ants.make_image( mask, temp['tSTD'] ) + tsnrmask = ants.threshold_image( tsnrmask, temp['threshold_std'], temp['tSTD'].max() ) + M = ants.timeseries_to_matrix( boldImage, tsnrmask ) + components = None + basis = np.array([]) + if filter_type in ('polynomial', False): + M, basis = regress_poly(degree, M) +# M = M / compute_tSTD(M, 1.)['tSTD'] + # "The covariance matrix C = MMT was constructed and decomposed into its + # principal components using a singular value decomposition." + M = clean_data( M, standardize=True ) + u, _, _ = linalg.svd(M, full_matrices=False) + if components is None: + components = u[:, :ncompcor] + else: + components = np.hstack((components, u[:, :ncompcor])) + if components is None and ncompcor > 0: + raise ValueError('No components found') + return { 'components': components, 'basis': basis } diff --git a/MindEyeV2/antspy/ants/registration/__init__.py b/MindEyeV2/antspy/ants/registration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b44ef04c648819fe0767c002310593932db9b68a --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/__init__.py @@ -0,0 +1,15 @@ +from .affine_initializer import affine_initializer +from .apply_transforms import apply_transforms, apply_transforms_to_points +from .average_transform import average_affine_transform, average_affine_transform_no_rigid +from .build_template import build_template +from .compose_displacement_fields import compose_displacement_fields +from .create_jacobian_determinant_image import create_jacobian_determinant_image, deformation_gradient +from .create_warped_grid import create_warped_grid +from .fit_bspline_displacement_field import fit_bspline_displacement_field +from .fit_bspline_object_to_scattered_data import fit_bspline_object_to_scattered_data +from .fit_thin_plate_spline_displacement_field import fit_thin_plate_spline_displacement_field +from .integrate_velocity_field import integrate_velocity_field +from .invert_displacement_field import invert_displacement_field +from .landmark_transforms import fit_transform_to_paired_points, fit_time_varying_transform_to_point_sets +from .registration import registration, motion_correction, label_image_registration +from .simulate_displacement_field import simulate_displacement_field diff --git a/MindEyeV2/antspy/ants/registration/affine_initializer.py b/MindEyeV2/antspy/ants/registration/affine_initializer.py new file mode 100644 index 0000000000000000000000000000000000000000..7c9a36eea140c4233ab6df5652e3f22b8d57ecc6 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/affine_initializer.py @@ -0,0 +1,70 @@ + +__all__ = ['affine_initializer'] + +import warnings +from tempfile import mktemp + +from ants.internal import get_lib_fn, process_arguments + + +def affine_initializer(fixed_image, moving_image, search_factor=20, + radian_fraction=0.1, use_principal_axis=False, + local_search_iterations=10, mask=None, txfn=None ): + """ + A multi-start optimizer for affine registration + Searches over the sphere to find a good initialization for further + registration refinement, if needed. This is a wrapper for the ANTs + function antsAffineInitializer. + + ANTsR function: `affineInitializer` + + Arguments + --------- + fixed_image : ANTsImage + the fixed reference image + moving_image : ANTsImage + the moving image to be mapped to the fixed space + search_factor : scalar + degree of increments on the sphere to search + radian_fraction : scalar + between zero and one, defines the arc to search over + use_principal_axis : boolean + boolean to initialize by principal axis + local_search_iterations : scalar + gradient descent iterations + mask : ANTsImage (optional) + optional mask to restrict registration + txfn : string (optional) + filename for the transformation + + Returns + ------- + ndarray + transformation matrix + + Example + ------- + >>> import ants + >>> fi = ants.image_read(ants.get_ants_data('r16')) + >>> mi = ants.image_read(ants.get_ants_data('r27')) + >>> txfile = ants.affine_initializer( fi, mi ) + >>> tx = ants.read_transform(txfile, dimension=2) + """ + + if txfn is None: + txfn = mktemp(suffix='.mat') + + veccer = [fixed_image.dimension, fixed_image, moving_image, txfn, + search_factor, radian_fraction, int(use_principal_axis), + local_search_iterations] + if mask is not None: + veccer.append(mask) + + xxx = process_arguments(veccer) + libfn = get_lib_fn('antsAffineInitializer') + retval = libfn(xxx) + + if retval != 0: + warnings.warn('ERROR: Non-zero exit status!') + + return txfn \ No newline at end of file diff --git a/MindEyeV2/antspy/ants/registration/apply_transforms.py b/MindEyeV2/antspy/ants/registration/apply_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..2914af7ef4a3e00df7fa237c6f93db63cc316204 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/apply_transforms.py @@ -0,0 +1,316 @@ + + +__all__ = ['apply_transforms', + 'apply_transforms_to_points'] + +import os + +import ants +from ants.internal import get_lib_fn, process_arguments + + +def apply_transforms(fixed, moving, transformlist, + interpolator='linear', imagetype=0, + whichtoinvert=None, compose=None, + defaultvalue=0, singleprecision=False, verbose=False, **kwargs): + """ + Apply a transform list to map an image from one domain to another. + In image registration, one computes mappings between (usually) pairs + of images. These transforms are often a sequence of increasingly + complex maps, e.g. from translation, to rigid, to affine to deformation. + The list of such transforms is passed to this function to interpolate one + image domain into the next image domain, as below. The order matters + strongly and the user is advised to familiarize with the standards + established in examples. + + ANTsR function: `antsApplyTransforms` + + Arguments + --------- + fixed : ANTsImage + fixed image defining domain into which the moving image is transformed. The output will + have the same pixel type as this image. + + moving : AntsImage + moving image to be mapped to fixed space. + + transformlist : list of strings + list of transforms generated by ants.registration where each transform is a filename. + + interpolator : string + Choice of interpolator. Supports partial matching. + linear + nearestNeighbor + multiLabel for label images (deprecated, prefer genericLabel) + gaussian + bSpline + cosineWindowedSinc + welchWindowedSinc + hammingWindowedSinc + lanczosWindowedSinc + genericLabel use this for label images + + imagetype : integer + choose 0/1/2/3 mapping to scalar/vector/tensor/time-series + + whichtoinvert : list of booleans (optional) + Must be same length as transformlist. + whichtoinvert[i] is True if transformlist[i] is a matrix, + and the matrix should be inverted. If transformlist[i] is a + warp field, whichtoinvert[i] must be False. + If the transform list is a matrix followed by a warp field, + whichtoinvert defaults to (True,False). Otherwise it defaults + to [False]*len(transformlist)). + + compose : string (optional) + if it is a string pointing to a valid file location, + this will force the function to return a composite transformation filename. + + defaultvalue : scalar + Default voxel value for mappings outside the image domain. + + singleprecision : boolean + if True, use float32 for computations. This is useful for reducing memory + usage for large datasets, at the cost of precision. + + verbose : boolean + print command and run verbose application of transform. + + kwargs : keyword arguments + extra parameters + + Returns + ------- + ANTsImage or string (transformation filename) + + Example + ------- + >>> import ants + >>> fixed = ants.image_read( ants.get_ants_data('r16') ) + >>> moving = ants.image_read( ants.get_ants_data('r64') ) + >>> fixed = ants.resample_image(fixed, (64,64), 1, 0) + >>> moving = ants.resample_image(moving, (64,64), 1, 0) + >>> mytx = ants.registration(fixed=fixed , moving=moving , + type_of_transform = 'SyN' ) + >>> mywarpedimage = ants.apply_transforms( fixed=fixed, moving=moving, + transformlist=mytx['fwdtransforms'] ) + """ + + if not isinstance(transformlist, (tuple, list)) and (transformlist is not None): + transformlist = [transformlist] + + accepted_interpolators = {"linear", "nearestNeighbor", "multiLabel", "gaussian", + "bSpline", "cosineWindowedSinc", "welchWindowedSinc", + "hammingWindowedSinc", "lanczosWindowedSinc", "genericLabel"} + + if interpolator not in accepted_interpolators: + raise ValueError('interpolator not supported - see %s' % accepted_interpolators) + + args = [fixed, moving, transformlist, interpolator] + + output_pixel_type = 'float' if singleprecision else 'double' + + if not isinstance(fixed, str): + if ants.is_image(fixed) and ants.is_image(moving): + for tl_path in transformlist: + if not os.path.exists(tl_path): + raise Exception('Transform %s does not exist' % tl_path) + + inpixeltype = fixed.pixeltype + fixed = fixed.clone(output_pixel_type) + moving = moving.clone(output_pixel_type) + warpedmovout = moving.clone(output_pixel_type) + f = fixed + m = moving + if (moving.dimension == 4) and (fixed.dimension == 3) and (imagetype == 0): + raise Exception('Set imagetype 3 to transform time series images.') + + wmo = warpedmovout + mytx = [] + if whichtoinvert is None or (isinstance(whichtoinvert, (tuple,list)) and (sum([w is not None for w in whichtoinvert])==0)): + if (len(transformlist) == 2) and ('.mat' in transformlist[0]) and ('.mat' not in transformlist[1]): + whichtoinvert = (True, False) + else: + whichtoinvert = tuple([False]*len(transformlist)) + + if len(whichtoinvert) != len(transformlist): + raise ValueError('Transform list and inversion list must be the same length') + + for i in range(len(transformlist)): + ismat = False + if '.mat' in transformlist[i]: + ismat = True + if whichtoinvert[i] and (not ismat): + raise ValueError('Cannot invert transform %i (%s) because it is not a matrix' % (i, transformlist[i])) + if whichtoinvert[i]: + mytx = mytx + ['-t', '[%s,1]' % (transformlist[i])] + else: + mytx = mytx + ['-t', transformlist[i]] + + if compose is None: + args = ['-d', fixed.dimension, + '-i', m, + '-o', wmo, + '-r', f, + '-n', interpolator] + args = args + mytx + if compose: + tfn = '%scomptx.nii.gz' % compose if not compose.endswith('.h5') else compose + else: + tfn = 'NA' + if compose is not None: + mycompo = '[%s,1]' % tfn + args = ['-d', fixed.dimension, + '-i', m, + '-o', mycompo, + '-r', f, + '-n', interpolator] + args = args + mytx + + myargs = process_arguments(args) + + myverb = int(verbose) + if verbose: + print(myargs) + + processed_args = myargs + ['-z', str(1), '-v', str(myverb), '--float', str(int(singleprecision)), '-e', str(imagetype), '-f', str(defaultvalue)] + libfn = get_lib_fn('antsApplyTransforms') + libfn(processed_args) + + if compose is None: + return warpedmovout.clone(inpixeltype) + else: + if os.path.exists(tfn): + return tfn + else: + return None + + else: + return 1 + else: + args = args + ['-z', str(1), '--float', str(int(singleprecision)), '-e', imagetype, '-f', defaultvalue] + processed_args = process_arguments(args) + libfn = get_lib_fn('antsApplyTransforms') + libfn(processed_args) + + + + + + +def apply_transforms_to_points( dim, points, transformlist, + whichtoinvert=None, verbose=False ): + """ + Apply a transform list to map a pointset from one domain to + another. In registration, one computes mappings between pairs of + domains. These transforms are often a sequence of increasingly + complex maps, e.g. from translation, to rigid, to affine to + deformation. The list of such transforms is passed to this + function to interpolate one image domain into the next image + domain, as below. The order matters strongly and the user is + advised to familiarize with the standards established in examples. + Importantly, point mapping goes the opposite direction of image + mapping, for both reasons of convention and engineering. + + ANTsR function: `antsApplyTransformsToPoints` + + Arguments + --------- + dim: integer + dimensionality of the transformation. + + points: data frame + moving point set with n-points in rows of at least dim + columns - we maintain extra information in additional + columns. this should be a data frame with columns names x, y, z, t. + + transformlist : list of strings + list of transforms generated by ants.registration where each transform is a filename. + + whichtoinvert : list of booleans (optional) + Must be same length as transformlist. + whichtoinvert[i] is True if transformlist[i] is a matrix, + and the matrix should be inverted. If transformlist[i] is a + warp field, whichtoinvert[i] must be False. + If the transform list is a matrix followed by a warp field, + whichtoinvert defaults to (True,False). Otherwise it defaults + to [False]*len(transformlist)). + + verbose : boolean + + Returns + ------- + data frame of transformed points + + Example + ------- + >>> import ants + >>> fixed = ants.image_read( ants.get_ants_data('r16') ) + >>> moving = ants.image_read( ants.get_ants_data('r27') ) + >>> reg = ants.registration( fixed, moving, 'Affine' ) + >>> d = {'x': [128, 127], 'y': [101, 111]} + >>> pts = pd.DataFrame(data=d) + >>> ptsw = ants.apply_transforms_to_points( 2, pts, reg['fwdtransforms']) + """ + + if not isinstance(transformlist, (tuple, list)) and (transformlist is not None): + transformlist = [transformlist] + + args = [dim, points, transformlist, whichtoinvert] + + for tl_path in transformlist: + if not os.path.exists(tl_path): + raise Exception('Transform %s does not exist' % tl_path) + + mytx = [] + + if whichtoinvert is None or (isinstance(whichtoinvert, (tuple,list)) and (sum([w is not None for w in whichtoinvert])==0)): + if (len(transformlist) == 2) and ('.mat' in transformlist[0]) and ('.mat' not in transformlist[1]): + whichtoinvert = (True, False) + else: + whichtoinvert = tuple([False]*len(transformlist)) + + if len(whichtoinvert) != len(transformlist): + raise ValueError('Transform list and inversion list must be the same length') + + for i in range(len(transformlist)): + ismat = False + if '.mat' in transformlist[i]: + ismat = True + if whichtoinvert[i] and (not ismat): + raise ValueError('Cannot invert transform %i (%s) because it is not a matrix' % (i, transformlist[i])) + if whichtoinvert[i]: + mytx = mytx + ['-t', '[%s,1]' % (transformlist[i])] + else: + mytx = mytx + ['-t', transformlist[i]] + if dim == 2: + pointsSub = points[['x','y']] + if dim == 3: + pointsSub = points[['x','y','z']] + if dim == 4: + pointsSub = points[['x','y','z','t']] + pointImage = ants.make_image( pointsSub.shape, pointsSub.values.flatten()) + pointsOut = pointImage.clone() + args = ['-d', dim, + '-i', pointImage, + '-o', pointsOut ] + args = args + mytx + myargs = process_arguments(args) + + myverb = int(verbose) + if verbose: + print(myargs) + + processed_args = myargs + [ '-f', str(1), '--precision', str(0)] + libfn = get_lib_fn('antsApplyTransformsToPoints') + libfn(processed_args) + mynp = pointsOut.numpy() + pointsOutDF = points.copy() + pointsOutDF['x'] = mynp[:,0] + if dim >= 2: + pointsOutDF['y'] = mynp[:,1] + if dim >= 3: + pointsOutDF['z'] = mynp[:,2] + if dim >= 4: + pointsOutDF['t'] = mynp[:,3] + return pointsOutDF diff --git a/MindEyeV2/antspy/ants/registration/average_transform.py b/MindEyeV2/antspy/ants/registration/average_transform.py new file mode 100644 index 0000000000000000000000000000000000000000..e12a93ca7d8fc964f60f4a83beabc6b25997abea --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/average_transform.py @@ -0,0 +1,48 @@ + +from tempfile import mktemp +import os + +import ants +from ants.internal import get_lib_fn, process_arguments + +__all__ = ['average_affine_transform', + 'average_affine_transform_no_rigid'] + + +def _average_affine_transform_driver(transformlist, referencetransform=None, funcname="AverageAffineTransform"): + """ + takes a list of transforms (files at the moment) + and returns the average + """ + + # AverageAffineTransform deals with transform files, + # so this function will need to deal with already + # loaded files. Doesn't look like the magic + # available for images has been added for transforms. + res_temp_file = mktemp(suffix='.mat') + + # could do some stuff here to cope with transform lists that + # aren't files + + # load one of the transforms to figure out the dimension + tf = ants.read_transform(transformlist[0]) + if referencetransform is None: + args = [tf.dimension, res_temp_file] + transformlist + else: + args = [tf.dimension, res_temp_file] + ['-R', referencetransform] + transformlist + pargs = process_arguments(args) + print(pargs) + libfun = get_lib_fn(funcname) + status = libfun(pargs) + + res = ants.read_transform(res_temp_file) + os.remove(res_temp_file) + return res + +def average_affine_transform(transformlist, referencetransform=None): + return _average_affine_transform_driver(transformlist, referencetransform, "AverageAffineTransform") + + +def average_affine_transform_no_rigid(transformlist, referencetransform=None): + return _average_affine_transform_driver(transformlist, referencetransform, "AverageAffineTransformNoRigid") + diff --git a/MindEyeV2/antspy/ants/registration/build_template.py b/MindEyeV2/antspy/ants/registration/build_template.py new file mode 100644 index 0000000000000000000000000000000000000000..4b089637b8f094424b208b2c0a66f19ff55cc306 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/build_template.py @@ -0,0 +1,137 @@ +__all__ = ["build_template"] + +import numpy as np +import os +import shutil +from tempfile import mktemp + +import ants + +def build_template( + initial_template=None, + image_list=None, + iterations=3, + gradient_step=0.2, + blending_weight=0.75, + weights=None, + useNoRigid=True, + output_dir=None, + **kwargs +): + """ + Estimate an optimal template from an input image_list + + ANTsR function: N/A + + Arguments + --------- + initial_template : ANTsImage + initialization for the template building + + image_list : ANTsImages + images from which to estimate template + + iterations : integer + number of template building iterations + + gradient_step : scalar + for shape update gradient + + blending_weight : scalar + weight for image blending + + weights : vector + weight for each input image + + useNoRigid : boolean + equivalent of -y in the script. Template update + step will not use the rigid component if this is True. + + output_dir : path + directory name where intermediate transforms are written + + kwargs : keyword args + extra arguments passed to ants registration + + Returns + ------- + ANTsImage + + Example + ------- + >>> import ants + >>> image = ants.image_read( ants.get_ants_data('r16') ) + >>> image2 = ants.image_read( ants.get_ants_data('r27') ) + >>> image3 = ants.image_read( ants.get_ants_data('r85') ) + >>> timage = ants.build_template( image_list = ( image, image2, image3 ) ).resample_image( (45,45)) + >>> timagew = ants.build_template( image_list = ( image, image2, image3 ), weights = (5,1,1) ) + """ + work_dir = mktemp() if output_dir is None else output_dir + + def make_outprefix(k: int): + os.makedirs(os.path.join(work_dir, f"img{k:04d}"), exist_ok=True) + return os.path.join(work_dir, f"img{k:04d}", "out") + + if "type_of_transform" not in kwargs: + type_of_transform = "SyN" + else: + type_of_transform = kwargs.pop("type_of_transform") + + if weights is None: + weights = np.repeat(1.0 / len(image_list), len(image_list)) + weights = [x / sum(weights) for x in weights] + if initial_template is None: + initial_template = image_list[0] * 0 + for i in range(len(image_list)): + temp = image_list[i] * weights[i] + temp = ants.resample_image_to_target(temp, initial_template) + initial_template = initial_template + temp + + xavg = initial_template.clone() + for i in range(iterations): + affinelist = [] + for k in range(len(image_list)): + w1 = ants.registration( + xavg, image_list[k], type_of_transform=type_of_transform, outprefix=make_outprefix(k), **kwargs + ) + L = len(w1["fwdtransforms"]) + # affine is the last one + affinelist.append(w1["fwdtransforms"][L-1]) + + if k == 0: + if L == 2: + wavg = ants.image_read(w1["fwdtransforms"][0]) * weights[k] + xavgNew = w1["warpedmovout"] * weights[k] + else: + if L == 2: + wavg = wavg + ants.image_read(w1["fwdtransforms"][0]) * weights[k] + xavgNew = xavgNew + w1["warpedmovout"] * weights[k] + + if useNoRigid: + avgaffine = ants.average_affine_transform_no_rigid(affinelist) + else: + avgaffine = ants.average_affine_transform(affinelist) + afffn = os.path.join(work_dir, "avgAffine.mat") + ants.write_transform(avgaffine, afffn) + + if L == 2: + print(wavg.abs().mean()) + wscl = (-1.0) * gradient_step + wavg = wavg * wscl + # apply affine to the nonlinear? + # need to save the average + wavgA = ants.apply_transforms(fixed=xavgNew, moving=wavg, imagetype=1, transformlist=afffn, whichtoinvert=[1]) + wavgfn = os.path.join(work_dir, "avgWarp.nii.gz") + ants.image_write(wavgA, wavgfn) + xavg = ants.apply_transforms(fixed=xavgNew, moving=xavgNew, transformlist=[wavgfn, afffn], whichtoinvert=[0, 1]) + else: + xavg = ants.apply_transforms(fixed=xavgNew, moving=xavgNew, transformlist=[afffn], whichtoinvert=[1]) + + if blending_weight is not None: + xavg = xavg * blending_weight + ants.iMath(xavg, "Sharpen") * ( + 1.0 - blending_weight + ) + + if output_dir is None: + shutil.rmtree(work_dir) + return xavg diff --git a/MindEyeV2/antspy/ants/registration/compose_displacement_fields.py b/MindEyeV2/antspy/ants/registration/compose_displacement_fields.py new file mode 100644 index 0000000000000000000000000000000000000000..9009e4789699fa36966281523b2d2bd9231525fb --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/compose_displacement_fields.py @@ -0,0 +1,33 @@ + +__all__ = ['compose_displacement_fields'] + +import ants +from ants.internal import get_lib_fn + + +def compose_displacement_fields(displacement_field, + warping_field): + """ + Compose displacement fields. + + Arguments + --------- + displacement_field : ANTsImage displacement field + displacement field + + warping_field : ANTsImage displacement field + warping field + + + Example + ------- + >>> import ants + """ + + libfn = get_lib_fn('composeDisplacementFieldsD%i' % displacement_field.dimension) + comp_field = libfn(displacement_field.pointer, warping_field.pointer) + + new_image = ants.from_pointer(comp_field).clone('float') + return new_image + + diff --git a/MindEyeV2/antspy/ants/registration/create_jacobian_determinant_image.py b/MindEyeV2/antspy/ants/registration/create_jacobian_determinant_image.py new file mode 100644 index 0000000000000000000000000000000000000000..e566129e32c9050ae90b978893a9c99b8cea4eb8 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/create_jacobian_determinant_image.py @@ -0,0 +1,175 @@ + + + +__all__ = ['create_jacobian_determinant_image', + 'deformation_gradient'] + +from tempfile import mktemp + +import ants +from ants.internal import get_lib_fn, process_arguments + + +def deformation_gradient( warp_image, to_rotation=False, py_based=False ): + """ + Compute the deformation gradient from an image containing a warp (deformation) + + ANTsR function: `NA` + + Arguments + --------- + warp_image : ANTsImage (or filename if not py_based) + image that defines the deformation field (vector pixels) + + to_rotation : boolean maps deformation gradient to a rotation matrix + + py_based: boolean uses pure python implementation (maybe slow) + + Returns + ------- + ANTsImage with dimension*dimension components indexed in order U_xyz, V_xyz, W_xyz + where U is the x-component of deformation and xyz are spatial. + + Note + ------- + the to_rotation option is still experimental. use with caution. + + Example + ------- + >>> import ants + >>> fi = ants.image_read( ants.get_ants_data('r16')) + >>> mi = ants.image_read( ants.get_ants_data('r64')) + >>> fi = ants.resample_image(fi,(128,128),1,0) + >>> mi = ants.resample_image(mi,(128,128),1,0) + >>> mytx = ants.registration(fixed=fi , moving=mi, type_of_transform = ('SyN') ) + >>> dg = ants.deformation_gradient( ants.image_read( mytx['fwdtransforms'][0] ) ) + """ + import numpy as np + def polar_decomposition(X): + U, d, V = np.linalg.svd(X, full_matrices=False) + P = np.matmul(U, np.matmul(np.diag(d), np.transpose(U))) + Z = np.matmul(U, V) + if np.linalg.det(Z) < 0: + n = X.shape[0] + reflection_matrix = np.identity(n) + reflection_matrix[0,0] = -1.0 + Z = np.matmul(Z, reflection_matrix) + return({"P" : P, "Z" : Z, "Xtilde" : np.matmul(P, Z)}) + if not py_based: + if ants.is_image(warp_image): + txuse = mktemp(suffix='.nii.gz') + ants.image_write(warp_image, txuse) + else: + txuse = warp_image + warp_image=ants.image_read(txuse) + if not ants.is_image(warp_image): + raise RuntimeError("antsimage is required") + writtenimage = mktemp(suffix='.nrrd') + dimage = warp_image.split_channels()[0].clone('double') + dim = dimage.dimension + tshp = dimage.shape + args2 = [dim, txuse, writtenimage, int(0), int(0), int(1)] + processed_args = process_arguments(args2) + libfn = get_lib_fn('CreateJacobianDeterminantImage') + libfn(processed_args) + dg = ants.image_read(writtenimage) + if to_rotation: + newshape = tshp + (dim,dim) + dg = np.reshape( dg.numpy(), newshape ) + it=np.ndindex(tshp) + for i in it: + dg[i]=polar_decomposition( dg[i] )['Z'] + newshape = tshp + (dim*dim,) + dg = np.reshape( dg, newshape ) + dg = ants.from_numpy( dg, has_components=True ) + dg = ants.copy_image_info( dimage, dg ) + import os + os.remove( writtenimage ) + return dg + if py_based: + if not ants.is_image(warp_image): + raise RuntimeError("antsimage is required") + dim = warp_image.dimension + warpnp=warp_image.numpy() + tshp=warp_image.shape + tdir=warp_image.direction + spc = warp_image.spacing + it=np.ndindex(tshp) + # print("first we need to rotate the warp by the direction cosines") + for i in it: + warpnp[i]=np.dot( tdir,warpnp[i]) + # print("second get deformation gradient") + dg = [] + for k in range(dim): + if dim == 2: + temp=np.stack( np.gradient( warpnp[...,k], spc[0], spc[1], axis=range(dim) ), axis=dim) + if dim == 3: + temp=np.stack( np.gradient( warpnp[...,k], spc[0], spc[1], spc[2], axis=range(dim) ), axis=dim) + dg.append(temp) + dg = np.stack(dg,axis=dim+1) + it=np.ndindex(tshp) + ident = np.eye( dim ) + for i in it: + dg[i]=dg[i]+ident + if to_rotation: + it=np.ndindex(tshp) + for i in it: + dg[i]=polar_decomposition( dg[i] )['Z'] + newshape = tshp + (dim*dim,) + dg = np.reshape( dg, newshape ) + dg = ants.from_numpy( dg, has_components=True ) + dg = ants.copy_image_info( warp_image, dg ) + return dg + + + +def create_jacobian_determinant_image(domain_image, tx, do_log=False, geom=False): + """ + Compute the jacobian determinant from a transformation file + + ANTsR function: `createJacobianDeterminantImage` + + Arguments + --------- + domain_image : ANTsImage + image that defines transformation domain + + tx : string + deformation transformation file name + + do_log : boolean + return the log jacobian + + geom : bolean + use the geometric jacobian calculation (boolean) + + Returns + ------- + ANTsImage + + Example + ------- + >>> import ants + >>> fi = ants.image_read( ants.get_ants_data('r16')) + >>> mi = ants.image_read( ants.get_ants_data('r64')) + >>> fi = ants.resample_image(fi,(128,128),1,0) + >>> mi = ants.resample_image(mi,(128,128),1,0) + >>> mytx = ants.registration(fixed=fi , moving=mi, type_of_transform = ('SyN') ) + >>> jac = ants.create_jacobian_determinant_image(fi,mytx['fwdtransforms'][0],1) + """ + dim = domain_image.dimension + if ants.is_image(tx): + txuse = mktemp(suffix='.nii.gz') + ants.image_write(tx, txuse) + else: + txuse = tx + #args = [dim, txuse, do_log] + dimage = domain_image.clone('double') + args2 = [dim, txuse, dimage, int(do_log), int(geom)] + processed_args = process_arguments(args2) + libfn = get_lib_fn('CreateJacobianDeterminantImage') + libfn(processed_args) + jimage = args2[2].clone('float') + + return jimage + diff --git a/MindEyeV2/antspy/ants/registration/create_warped_grid.py b/MindEyeV2/antspy/ants/registration/create_warped_grid.py new file mode 100644 index 0000000000000000000000000000000000000000..f2241bff6c702cd6db3fd5c1e8abcae5798a05a5 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/create_warped_grid.py @@ -0,0 +1,105 @@ + + +__all__ = ['create_warped_grid'] + +import numpy as np + +import ants + + +def create_warped_grid(image, grid_step=10, grid_width=2, grid_directions=(True, True), + fixed_reference_image=None, transform=None, foreground=1, background=0): + """ + Deforming a grid is a helpful way to visualize a deformation field. + This function enables a user to define the grid parameters + and apply a deformable map to that grid. + + ANTsR function: `createWarpedGrid` + + Arguments + --------- + image : ANTsImage + input image + + grid_step : scalar + width of grid blocks + + grid_width : scalar + width of grid lines + + grid_directions : tuple of booleans + directions in which to draw grid lines, boolean vector + + fixed_reference_image : ANTsImage (optional) + reference image space + + transform : list/tuple of strings (optional) + vector of transforms + + foreground : scalar + intensity value for grid blocks + + background : scalar + intensity value for grid lines + + Returns + ------- + ANTsImage + + Example + ------- + >>> import ants + >>> fi = ants.image_read( ants.get_ants_data( 'r16' ) ) + >>> mi = ants.image_read( ants.get_ants_data( 'r64' ) ) + >>> mygr = ants.create_warped_grid( mi ) + >>> mytx = ants.registration(fixed=fi, moving=mi, type_of_transform = ('SyN') ) + >>> mywarpedgrid = ants.create_warped_grid( mygr, grid_directions=(False,True), + transform=mytx['fwdtransforms'], fixed_reference_image=fi ) + """ + if ants.is_image(image): + if len(grid_directions) != image.dimension: + grid_directions = [True]*image.dimension + garr = image.numpy() * 0 + foreground + else: + if not isinstance(image, (list, tuple)): + raise ValueError('image arg must be ANTsImage or list or tuple') + if len(grid_directions) != len(image): + grid_directions = [True]*len(image) + garr = np.zeros(image) + foreground + image = ants.from_numpy(garr) + + idim = garr.ndim + gridw = grid_width + + for d in range(idim): + togrid = np.arange(-1, garr.shape[d]-1, step=grid_step) + for i in range(len(togrid)): + if (d == 0) & (idim == 3) & (grid_directions[d]): + garr[togrid[i]:(togrid[i]+gridw),...] = background + garr[0,...] = background + garr[-1,...] = background + if (d == 1) & (idim == 3) & (grid_directions[d]): + garr[:,togrid[i]:(togrid[i]+gridw),:] = background + garr[:,0,:] = background + garr[:,-1,:] = background + if (d == 2) & (idim == 3) & (grid_directions[d]): + garr[...,togrid[i]:(togrid[i]+gridw)] = background + garr[...,0] = background + garr[...,-1] = background + if (d == 0) & (idim == 2) & (grid_directions[d]): + garr[togrid[i]:(togrid[i]+gridw),:] = background + garr[0,:] = background + garr[-1,:] = background + if (d == 1) & (idim == 2) & (grid_directions[d]): + garr[:,togrid[i]:(togrid[i]+gridw)] = background + garr[:,0] = background + garr[:,-1] = background + + + gimage = image.new_image_like(garr) + + if (transform is not None) and (fixed_reference_image is not None): + return ants.apply_transforms( fixed=fixed_reference_image, moving=gimage, + transformlist=transform ) + else: + return gimage diff --git a/MindEyeV2/antspy/ants/registration/fit_bspline_displacement_field.py b/MindEyeV2/antspy/ants/registration/fit_bspline_displacement_field.py new file mode 100644 index 0000000000000000000000000000000000000000..323591c61aa7cf8252d7187579e50c40dcdac1ef --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/fit_bspline_displacement_field.py @@ -0,0 +1,202 @@ +__all__ = ["fit_bspline_displacement_field"] + +import numpy as np + +import ants +from ants.internal import get_lib_fn + + +def fit_bspline_displacement_field(displacement_field=None, + displacement_weight_image=None, + displacement_origins=None, + displacements=None, + displacement_weights=None, + origin=None, + spacing=None, + size=None, + direction=None, + number_of_fitting_levels=4, + mesh_size=1, + spline_order=3, + enforce_stationary_boundary=True, + estimate_inverse=False, + rasterize_points=False): + + """ + Fit a b-spline object to a dense displacement field image and/or a set of points + with associated displacements and smooths them using B-splines. The inverse + can also be estimated.. This is basically a wrapper for the ITK filter + + https://itk.org/Doxygen/html/classitk_1_1DisplacementFieldToBSplineImageFilter.html} + + which, in turn is a wrapper for the ITK filter used for the function + fit_bspline_object_to_scattered_data. + + ANTsR function: `fitBsplineToDisplacementField` + + Arguments + --------- + displacement_field : ANTs image + Input displacement field. Either this and/or the points must be specified. + + displacement_weight_image : ANTs image + Input image defining weighting of the voxelwise displacements in the displacement_field. I + If None, defaults to identity weighting for each displacement. Default = None. + + displacement_origins : 2-D numpy array + Matrix (number_of_points x dimension) defining the origins of the input + displacement points. Default = None. + + displacements : 2-D numpy array + Matrix (number_of_points x dimension) defining the displacements of the input + displacement points. Default = None. + + displacement_weights : 1-D numpy array + Array defining the individual weighting of the corresponding scattered data value. + Default = None meaning all values are weighted the same. + + origin : n-D tuple + Defines the physical origin of the B-spline object. + + spacing : n-D tuple + Defines the physical spacing of the B-spline object. + + size : n-D tuple + Defines the size (length) of the B-spline object. Note that the length of the + B-spline object in dimension d is defined as + spacing[d] * size[d]-1. + + direction : 2-D numpy array + Booleans defining whether or not the corresponding parametric dimension is + closed (e.g., closed loop). Default = None. + + number_of_fitting_levels : integer + Specifies the number of fitting levels. + + mesh_size : n-D tuple + Defines the mesh size at the initial fitting level. + + spline_order : integer + Spline order of the B-spline object. Default = 3. + + enforce_stationary_boundary : boolean + Ensure no displacements on the image boundary. Default = True. + + estimate_inverse : boolean + Estimate the inverse displacement field. Default = False. + + rasterize_points : boolean + Use nearest neighbor rasterization of points for estimating the + field (potential speed-up). Default = False. + + Returns + ------- + Returns an ANTsImage. + + Example + ------- + >>> import ants + >>> import numpy as np + >>> points = np.array([[-50, -50]]) + >>> deltas = np.array([[10, 10]]) + >>> bspline_field = ants.fit_bspline_displacement_field( + >>> displacement_origins=points, displacements=deltas, + >>> origin=[0.0, 0.0], spacing=[1.0, 1.0], size=[100, 100], + >>> direction=np.array([[-1, 0], [0, -1]]), + >>> number_of_fitting_levels=4, mesh_size=(1, 1)) + """ + + if displacement_field is None and (displacement_origins is None or displacements is None): + raise ValueError("Missing input. Either a displacement field or input point set (origins + displacements) needs to be specified.") + + if displacement_field is None: + if origin is None or spacing is None or size is None or direction is None: + raise ValueError("If the displacement field is not specified, one must fully specify the input physical domain.") + + if displacement_field is not None and displacement_weight_image is None: + displacement_weight_image = ants.make_image(displacement_field.shape, voxval=1, + spacing=displacement_field.spacing, origin=displacement_field.origin, + direction=displacement_field.direction, has_components=False, pixeltype='float') + + if displacement_field is not None: + if origin is None: + origin = displacement_field.origin + if spacing is None: + spacing = displacement_field.spacing + if direction is None: + direction = displacement_field.direction + if size is None: + size = displacement_field.shape + + dimensionality = None + if displacement_field is not None: + dimensionality = displacement_field.dimension + else: + dimensionality = displacement_origins.shape[1] + if displacements.shape[1] != dimensionality: + raise ValueError("Dimensionality between origins and displacements does not match.") + + if displacement_origins is not None: + if displacement_weights is not None and (len(displacement_weights) != displacement_origins.shape[0]): + raise ValueError("Length of displacement weights must match the number of displacement points.") + else: + displacement_weights = np.ones(displacement_origins.shape[0]) + + if isinstance(mesh_size, int) == False and len(mesh_size) != dimensionality: + raise ValueError("Incorrect specification for mesh_size.") + + if origin is not None and len(origin) != dimensionality: + raise ValueError("Origin is not of length dimensionality.") + + if spacing is not None and len(spacing) != dimensionality: + raise ValueError("Spacing is not of length dimensionality.") + + if size is not None and len(size) != dimensionality: + raise ValueError("Size is not of length dimensionality.") + + if direction is not None and (direction.shape[0] != dimensionality and direction.shape[1] != dimensionality): + raise ValueError("Direction is not of shape dimensionality x dimensionality.") + + # It would seem that pybind11 doesn't really play nicely when the + # arguments are 'None' + + if origin is None: + origin = np.empty(0) + + if spacing is None: + spacing = np.empty(0) + + if size is None: + size = np.empty(0) + + if direction is None: + direction = np.empty((0, 0)) + + if displacement_origins is None: + displacement_origins = np.empty((0, 0)) + displacement_weights = np.empty(0) + else: + if displacement_weights is None: + displacement_weights = np.repeat(1.0, displacement_origins.shape[0]) + + number_of_control_points = list(np.array(mesh_size) + np.repeat(spline_order, dimensionality)) + + bspline_field = None + if displacement_field is not None: + libfn = get_lib_fn("fitBsplineDisplacementFieldD%i" % (dimensionality)) + bspline_field = libfn(displacement_field.pointer, displacement_weight_image.pointer, + displacement_origins, displacements, displacement_weights, + origin, spacing, size, direction, + number_of_fitting_levels, number_of_control_points, spline_order, + enforce_stationary_boundary, estimate_inverse) + elif displacement_field is None and displacements is not None: + libfn = get_lib_fn("fitBsplineDisplacementFieldToScatteredDataD%i" % (dimensionality)) + bspline_field = libfn(displacement_origins, displacements, displacement_weights, + origin, spacing, size, direction, + number_of_fitting_levels, number_of_control_points, spline_order, + enforce_stationary_boundary, estimate_inverse, rasterize_points) + + + bspline_displacement_field = ants.from_pointer(bspline_field).clone('float') + return bspline_displacement_field + diff --git a/MindEyeV2/antspy/ants/registration/fit_bspline_object_to_scattered_data.py b/MindEyeV2/antspy/ants/registration/fit_bspline_object_to_scattered_data.py new file mode 100644 index 0000000000000000000000000000000000000000..3ab5d615515157b67b363190248239bdbdef0218 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/fit_bspline_object_to_scattered_data.py @@ -0,0 +1,191 @@ +__all__ = ["fit_bspline_object_to_scattered_data"] + +import numpy as np + +import ants +from ants.internal import get_lib_fn + + +def fit_bspline_object_to_scattered_data(scattered_data, + parametric_data, + parametric_domain_origin, + parametric_domain_spacing, + parametric_domain_size, + is_parametric_dimension_closed=None, + data_weights=None, + number_of_fitting_levels=4, + mesh_size=1, + spline_order=3): + + """ + Fit a b-spline object to scattered data. This is basically a wrapper + for the ITK filter + + https://itk.org/Doxygen/html/classitk_1_1BSplineScatteredDataPointSetToImageFilter.html + + This filter is flexible in the possible objects that can be approximated. + Possibilities include: + + * 1/2/3/4-D curve + * 2-D surface in 3-D space (not available/templated) + * 2/3/4-D scalar field + * 2/3-D displacement field + * 2/3-D time-varying velocity field + + In order to understand the input parameters, it is important to understand + the difference between the parametric and data dimensions. A curve as one + parametric dimension but the data dimension can be 1-D, 2-D, 3-D, or 4-D. + In contrast, a 3-D displacement field has a parametric and data dimension + of 3. The scattered data is what's approximated by the B-spline object and + the parametric point is the location of scattered data within the domain of + the B-spline object. + + ANTsR function: `fitBsplineObjectToScatteredData` + + Arguments + --------- + scattered_data : 2-D numpy array + Defines the scattered data input to be approximated. Data is organized + by row --> data v, column ---> data dimension. + + parametric_data : 2-D numpy array + Defines the parametric location of the scattered data. Data is organized + by row --> parametric point, column --> parametric dimension. Note that + each row corresponds to the same row in the scatteredData. + + data_weights : 1-D numpy array + Defines the individual weighting of the corresponding scattered data value. + Default = None meaning all values are weighted the same. + + parametric_domain_origin : n-D tuple + Defines the parametric origin of the B-spline object. + + parametric_domain_spacing : n-D tuple + Defines the parametric spacing of the B-spline object. Defines the sampling + rate in the parametric domain. + + parametric_domain_size : n-D tuple + Defines the size (length) of the B-spline object. Note that the length of the + B-spline object in dimension d is defined as + parametric_domain_spacing[d] * parametric_domain_size[d]-1. + + is_parametric_dimension_closed : n-D tuple + Booleans defining whether or not the corresponding parametric dimension is + closed (e.g., closed loop). Default = None. + + number_of_fitting_levels : integer + Specifies the number of fitting levels. + + mesh_size : n-D tuple + Defines the mesh size at the initial fitting level. + + spline_order : integer + Spline order of the B-spline object. Default = 3. + + Returns + ------- + returns numpy array for B-spline curve (parametric dimension = 1). Otherwise, + returns an ANTsImage. + + Example + ------- + >>> # Perform 2-D curve example + >>> + >>> import ants, numpy + >>> import matplotlib.pyplot as plt + >>> x = numpy.linspace(-4, 4, num=100) + >>> y = numpy.exp(-numpy.multiply(x, x)) + numpy.random.uniform(-0.1, 0.1, len(x)) + >>> u = numpy.linspace(0, 1.0, num=len(x)) + >>> scattered_data = numpy.column_stack((x, y)) + >>> parametric_data = numpy.expand_dims(u, axis=-1) + >>> spacing = 1/(len(x)-1) * 1.0; + >>> bspline_curve = ants.fit_bspline_object_to_scattered_data(scattered_data, + >>> parametric_data, + >>> parametric_domain_origin=[0.0], parametric_domain_spacing=[spacing], + >>> parametric_domain_size=[len(x)], is_parametric_dimension_closed=None, + >>> number_of_fitting_levels=5, mesh_size=1) + >>> plt.plot(x, y, label='Noisy points') + >>> plt.plot(bspline_curve[:,0], bspline_curve[:,1], label='B-spline curve') + >>> plt.grid(True) + >>> plt.axis('tight') + >>> plt.legend(loc='upper left') + >>> plt.show() + >>> + >>> ########################################################################### + >>> + >>> # Perform 2-D scalar field (i.e., image) example + >>> + >>> import ants, numpy + >>> number_of_random_points = 10000 + >>> img = ants.image_read( ants.get_ants_data("r16")) + >>> img_array = img.numpy() + >>> row_indices = numpy.random.choice(range(2, img_array.shape[0]), number_of_random_points) + >>> col_indices = numpy.random.choice(range(2, img_array.shape[1]), number_of_random_points) + >>> scattered_data = numpy.zeros((number_of_random_points, 1)) + >>> parametric_data = numpy.zeros((number_of_random_points, 2)) + >>> for i in range(number_of_random_points): + >>> scattered_data[i,0] = img_array[row_indices[i], col_indices[i]] + >>> parametric_data[i,0] = row_indices[i] + >>> parametric_data[i,1] = col_indices[i] + >>> bspline_img = ants.fit_bspline_object_to_scattered_data( + >>> scattered_data, parametric_data, + >>> parametric_domain_origin=[0.0, 0.0], + >>> parametric_domain_spacing=[1.0, 1.0], + >>> parametric_domain_size = img.shape, + >>> number_of_fitting_levels=7, mesh_size=1) + >>> + >>> ants.plot(img, title="Original") + >>> ants.plot(bspline_img, title="B-spline approximation") + """ + + parametric_dimension = parametric_data.shape[1] + data_dimension = scattered_data.shape[1] + + if is_parametric_dimension_closed is None: + is_parametric_dimension_closed = np.repeat(False, parametric_dimension) + + if isinstance(mesh_size, int) == False and len(mesh_size) != parametric_dimension: + raise ValueError("Incorrect specification for mesh_size.") + + if len(parametric_domain_origin) != parametric_dimension: + raise ValueError("Origin is not of length parametric_dimension.") + + if len(parametric_domain_spacing) != parametric_dimension: + raise ValueError("Spacing is not of length parametric_dimension.") + + if len(parametric_domain_size) != parametric_dimension: + raise ValueError("Size is not of length parametric_dimension.") + + if len(is_parametric_dimension_closed) != parametric_dimension: + raise ValueError("Closed is not of length parametric_dimension.") + + number_of_control_points = mesh_size + spline_order + + if isinstance(number_of_control_points, int) == True: + number_of_control_points = np.repeat(number_of_control_points, parametric_dimension) + + if parametric_data.shape[0] != scattered_data.shape[0]: + raise ValueError("The number of points is not equal to the number of scattered data values.") + + if data_weights is None: + data_weights = np.repeat(1.0, parametric_data.shape[0]) + + if data_weights.ndim == 2: + data_weights = np.squeeze(data_weights) + + if len(data_weights) != parametric_data.shape[0]: + raise ValueError("The number of weights is not the same as the number of points.") + + libfn = get_lib_fn("fitBsplineObjectToScatteredDataP%iD%i" % (parametric_dimension, data_dimension)) + bspline_object = libfn(scattered_data.tolist(), parametric_data.tolist(), data_weights.tolist(), + parametric_domain_origin, parametric_domain_spacing, + parametric_domain_size, is_parametric_dimension_closed.tolist(), + number_of_fitting_levels, number_of_control_points.tolist(), + spline_order) + + if parametric_dimension == 1: + return np.array(bspline_object) + else: + bspline_image = ants.from_pointer(bspline_object).clone('float') + return bspline_image + diff --git a/MindEyeV2/antspy/ants/registration/fit_thin_plate_spline_displacement_field.py b/MindEyeV2/antspy/ants/registration/fit_thin_plate_spline_displacement_field.py new file mode 100644 index 0000000000000000000000000000000000000000..0e0ff773e1dedad3e42ca5d691dc8a81da9189ae --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/fit_thin_plate_spline_displacement_field.py @@ -0,0 +1,105 @@ +__all__ = ["fit_thin_plate_spline_displacement_field"] + +import numpy as np + +import ants +from ants.internal import get_lib_fn + + +def fit_thin_plate_spline_displacement_field(displacement_origins=None, + displacements=None, + origin=None, + spacing=None, + size=None, + direction=None): + + """ + Fit a thin-plate spline object to a a set of points with associated displacements. + This is basically a wrapper for the ITK filter + + https://itk.org/Doxygen/html/itkThinPlateSplineKernelTransform_8h.html + + ANTsR function: `fitThinPlateSplineToDisplacementField` + + Arguments + --------- + + displacement_origins : 2-D numpy array + Matrix (number_of_points x dimension) defining the origins of the input + displacement points. Default = None. + + displacements : 2-D numpy array + Matrix (number_of_points x dimension) defining the displacements of the input + displacement points. Default = None. + + origin : n-D tuple + Defines the physical origin of the B-spline object. + + spacing : n-D tuple + Defines the physical spacing of the B-spline object. + + size : n-D tuple + Defines the size (length) of the spline object. Note that the length of the + spline object in dimension d is defined as spacing[d] * size[d]-1. + + direction : 2-D numpy array + Booleans defining whether or not the corresponding parametric dimension is + closed (e.g., closed loop). Default = None. + + Returns + ------- + Returns an ANTsImage. + + Example + ------- + >>> import ants + >>> import numpy as np + >>> points = np.array([[-50, -50]]) + >>> deltas = np.array([[10, 10]]) + >>> tps_field = ants.fit_thin_plate_spline_displacement_field( + >>> displacement_origins=points, displacements=deltas, + >>> origin=[0.0, 0.0], spacing=[1.0, 1.0], size=[100, 100], + >>> direction=np.array([[-1, 0], [0, -1]])) + """ + + dimensionality = displacement_origins.shape[1] + if displacements.shape[1] != dimensionality: + raise ValueError("Dimensionality between origins and displacements does not match.") + + if displacement_origins is None or displacement_origins is None: + raise ValueError("Missing input. Input point set (origins + displacements) needs to be specified." ) + + if origin is not None and len(origin) != dimensionality: + raise ValueError("Origin is not of length dimensionality.") + + if spacing is not None and len(spacing) != dimensionality: + raise ValueError("Spacing is not of length dimensionality.") + + if size is not None and len(size) != dimensionality: + raise ValueError("Size is not of length dimensionality.") + + if direction is not None and (direction.shape[0] != dimensionality and direction.shape[1] != dimensionality): + raise ValueError("Direction is not of shape dimensionality x dimensionality.") + + # It would seem that pybind11 doesn't really play nicely when the + # arguments are 'None' + + if origin is None: + origin = np.empty(0) + + if spacing is None: + spacing = np.empty(0) + + if size is None: + size = np.empty(0) + + if direction is None: + direction = np.empty((0, 0)) + + tps_field = None + libfn = get_lib_fn("fitThinPlateSplineDisplacementFieldToScatteredDataD%i" % (dimensionality)) + tps_field = libfn(displacement_origins, displacements, origin, spacing, size, direction) + + tps_displacement_field = ants.from_pointer(tps_field).clone('float') + return tps_displacement_field + diff --git a/MindEyeV2/antspy/ants/registration/integrate_velocity_field.py b/MindEyeV2/antspy/ants/registration/integrate_velocity_field.py new file mode 100644 index 0000000000000000000000000000000000000000..9ea2386b1a5f728588fbca285247ed52c2b81cb4 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/integrate_velocity_field.py @@ -0,0 +1,48 @@ + +__all__ = ['integrate_velocity_field'] + +import ants +from ants.internal import get_lib_fn + + +def integrate_velocity_field(velocity_field, + lower_integration_bound=0.0, + upper_integration_bound=1.0, + number_of_integration_steps=10): + """ + Integrate velocity field. + + Arguments + --------- + velocity_field : ANTsImage velocity field + time-varying displacement field + + lower_integration_bound: float + Lower time bound for integration in [0, 1] + + upper_integration_bound: float + Upper time bound for integration in [0, 1] + + number_of_integation_steps: integer + Number of integration steps used in the Runge-Kutta solution + + Example + ------- + >>> import ants + >>> fi = ants.image_read( ants.get_data( "r16" ) ) + >>> mi = ants.image_read( ants.get_data( "r27" ) ) + >>> reg = ants.registration(fi, mi, "TV[2]") + >>> velocity_field = ants.image_read(reg['velocityfield'][0]) + >>> field = ants.integrate_velocity_field(velocity_field, 0.0, 1.0, 10) + >>> temp=ants.apply_ants_transform_to_image( + ants.transform_from_displacement_field( field ), mi, fi ) + """ + + libfn = get_lib_fn('integrateVelocityFieldD%i' % (velocity_field.dimension-1)) + integrated_field = libfn(velocity_field.pointer, lower_integration_bound, + upper_integration_bound, number_of_integration_steps) + + new_image = ants.from_pointer(integrated_field).clone('float') + return new_image + + diff --git a/MindEyeV2/antspy/ants/registration/invert_displacement_field.py b/MindEyeV2/antspy/ants/registration/invert_displacement_field.py new file mode 100644 index 0000000000000000000000000000000000000000..e0de961443f2786d19bd38f88573393fc247baf8 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/invert_displacement_field.py @@ -0,0 +1,51 @@ + +__all__ = ['invert_displacement_field'] + +import ants +from ants.internal import get_lib_fn + + +def invert_displacement_field(displacement_field, + inverse_field_initial_estimate, + maximum_number_of_iterations=20, + mean_error_tolerance_threshold=0.001, + max_error_tolerance_threshold=0.1, + enforce_boundary_condition=True): + """ + Invert displacement field. + + Arguments + --------- + displacement_field : ANTsImage displacement field + displacement field + + inverse_field_initial_estimate : ANTsImage displacement field + initial guess + + maximum_number_of_iterations : integer + number of iterations + + mean_error_tolerance_threshold : float + mean error tolerance threshold + + max_error_tolerance_threshold : float + max error tolerance threshold + + enforce_boundary_condition : bool + enforce stationary boundary condition + + + Example + ------- + >>> import ants + """ + + libfn = get_lib_fn('invertDisplacementFieldD%i' % displacement_field.dimension) + inverse_field = libfn(displacement_field.pointer, inverse_field_initial_estimate.pointer, + maximum_number_of_iterations, mean_error_tolerance_threshold, + max_error_tolerance_threshold, enforce_boundary_condition) + + new_image = ants.from_pointer(inverse_field).clone('float') + return new_image + + diff --git a/MindEyeV2/antspy/ants/registration/landmark_transforms.py b/MindEyeV2/antspy/ants/registration/landmark_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..41341e56e0bec02bb185d309c2b229b3a9eed263 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/landmark_transforms.py @@ -0,0 +1,843 @@ +__all__ = ["fit_transform_to_paired_points", + "fit_time_varying_transform_to_point_sets"] + +import numpy as np +import math +import time + +import ants + +def convergence_monitoring(values, window_size=10): + if len(values) >= window_size: + u = np.linspace(0.0, 1.0, num=window_size) + scattered_data = np.expand_dims(values[-window_size:], axis=-1) + parametric_data = np.expand_dims(u, axis=-1) + spacing = 1 / (window_size-1) + bspline_line = ants.fit_bspline_object_to_scattered_data(scattered_data, parametric_data, + parametric_domain_origin=[0.0], parametric_domain_spacing=[spacing], + parametric_domain_size=[window_size], number_of_fitting_levels=1, mesh_size=1, + spline_order=1) + bspline_slope = -(bspline_line[1][0] - bspline_line[0][0]) / spacing + return(bspline_slope) + else: + return None + + +def fit_transform_to_paired_points(moving_points, + fixed_points, + transform_type="affine", + regularization=1e-6, + domain_image=None, + number_of_fitting_levels=4, + mesh_size=1, + spline_order=3, + enforce_stationary_boundary=True, + displacement_weights=None, + number_of_compositions=10, + composition_step_size=0.5, + sigma=0.0, + convergence_threshold=1e-6, + number_of_time_steps=2, + number_of_integration_steps=100, + rasterize_points=False, + verbose=False + ): + """ + Estimate a transform from corresponding fixed and moving landmarks. + + ANTsR function: fitTransformToPairedPoints + + Arguments + --------- + moving_points : array + Moving points specified in physical space as a n x d matrix where n is the number + of points and d is the dimensionality. + + fixed_points : array + Fixed points specified in physical space as a n x d matrix where n is the number + of points and d is the dimensionality. + + transform_type : character + 'rigid', 'similarity', "affine', 'bspline', 'tps', 'diffeo', 'syn', or 'time-varying (tv)'. + + regularization : scalar + Ridge penalty in [0,1] for linear transforms. + + domain_image : ANTs image + Defines physical domain of the nonlinear transform. Must be defined for nonlinear + transforms. + + number_of_fitting_levels : integer + Integer specifying the number of fitting levels for the B-spline interpolation of the + displacement field. + + mesh_size : integer or array + Defines the mesh size at the initial fitting level for the B-spline interpolation of the + displacement field. + + spline_order : integer + Spline order of the B-spline displacement field. + + enforce_stationary_boundary : boolean + Ensure no displacements on the image boundary (B-spline only). + + displacement_weights : array + Defines the individual weighting of the corresponding scattered data value. Default = NULL + meaning all displacements are weighted the same. + + number_of_compositions : integer + Total number of compositions for the diffeomorphic transforms. + + composition_step_size : scalar + Scalar multiplication factor of the weighting of the update field for the diffeomorphic transforms. + + sigma : scalar + Gaussian smoothing standard deviation of the update field (in mm). + + convergence_threshold : scalar + Composition-based convergence parameter for the diff. transforms using a + window size of 10 values. + + number_of_time_steps : integer + Time-varying velocity field parameter. + + number_of_integration_steps : scalar + Number of steps used for integrating the velocity field. + + rasterize_points : boolean + Use nearest neighbor rasterization of points for estimating the update + field (potential speed-up). Default = False. + + verbose : bool + Print progress to the screen. + + Returns + ------- + + ANTs transform + + Example + ------- + >>> import ants + >>> import numpy as np + >>> fixed = np.array([[50.0,50.0],[200.0,50.0],[200.0,200.0]]) + >>> moving = np.array([[50.0,50.0],[50.0,200.0],[200.0,200.0]]) + >>> xfrm = ants.fit_transform_to_paired_points(moving, fixed, transform_type="affine") + >>> xfrm = ants.fit_transform_to_paired_points(moving, fixed, transform_type="rigid") + >>> xfrm = ants.fit_transform_to_paired_points(moving, fixed, transform_type="similarity") + >>> domain_image = ants.image_read(ants.get_ants_data("r16")) + >>> xfrm = ants.fit_transform_to_paired_points(moving, fixed, transform_type="bspline", domain_image=domain_image, number_of_fitting_levels=5) + >>> xfrm = ants.fit_transform_to_paired_points(moving, fixed, transform_type="diffeo", domain_image=domain_image, number_of_fitting_levels=6) + """ + + def polar_decomposition(X): + U, d, V = np.linalg.svd(X, full_matrices=False) + P = np.matmul(U, np.matmul(np.diag(d), np.transpose(U))) + Z = np.matmul(U, V) + if np.linalg.det(Z) < 0: + n = X.shape[0] + reflection_matrix = np.identity(n) + reflection_matrix[0,0] = -1.0 + Z = np.matmul(Z, reflection_matrix) + return({"P" : P, "Z" : Z, "Xtilde" : np.matmul(P, Z)}) + + def create_zero_displacement_field(domain_image): + field_array = np.zeros((*domain_image.shape, domain_image.dimension)) + field = ants.from_numpy(field_array, origin=domain_image.origin, + spacing=domain_image.spacing, direction=domain_image.direction, + has_components=True) + return(field) + + def create_zero_velocity_field(domain_image, number_of_time_points=2): + field_array = np.zeros((*domain_image.shape, number_of_time_points, domain_image.dimension)) + origin = (*domain_image.origin, 0.0) + spacing = (*domain_image.spacing, 1.0) + direction = np.eye(domain_image.dimension + 1) + direction[0:domain_image.dimension,0:domain_image.dimension] = domain_image.direction + field = ants.from_numpy(field_array, origin=origin, spacing=spacing, direction=direction, + has_components=True) + return(field) + + allowed_transforms = ['rigid', 'affine', 'similarity', 'bspline', 'tps', 'diffeo', 'syn', 'tv', 'time-varying'] + if not transform_type.lower() in allowed_transforms: + raise ValueError(transform_type + " transform not supported.") + + transform_type = transform_type.lower() + + if domain_image is None and transform_type in ['bspline', 'tps', 'diffeo', 'syn', 'tv', 'time-varying']: + raise ValueError("Domain image needs to be specified.") + + if not fixed_points.shape == moving_points.shape: + raise ValueError("Mismatch in the size of the point sets.") + + if regularization > 1: + regularization = 1 + elif regularization < 0: + regularization = 0 + + number_of_points = fixed_points.shape[0] + dimensionality = fixed_points.shape[1] + + if transform_type in ['rigid', 'affine', 'similarity']: + center_fixed = fixed_points.mean(axis=0) + center_moving = moving_points.mean(axis=0) + + x = fixed_points - center_fixed + y = moving_points - center_moving + + y_prior = np.concatenate((y, np.ones((number_of_points, 1))), axis=1) + + x11 = np.concatenate((x, np.ones((number_of_points, 1))), axis=1) + M = x11 * (1.0 - regularization) + regularization * y_prior + Minv = np.linalg.lstsq(M, y, rcond=None)[0] + + p = polar_decomposition(Minv[0:dimensionality, 0:dimensionality].T) + A = p['Xtilde'] + translation = Minv[dimensionality,:] + center_moving - center_fixed + + if transform_type in ['rigid', 'similarity']: + # Kabsch algorithm + # http://web.stanford.edu/class/cs273/refs/umeyama.pdf + + C = np.dot(y.T, x) + x_svd = np.linalg.svd(C * (1.0 - regularization) + np.eye(dimensionality) * regularization) + x_det = np.linalg.det(np.dot(x_svd[0], x_svd[2])) + + if x_det < 0: + x_svd[2][dimensionality-1, :] *= -1 + + A = np.dot(x_svd[0], x_svd[2]) + + if transform_type == 'similarity': + scaling = (math.sqrt((np.power(y, 2).sum(axis=1) / number_of_points).mean()) / + math.sqrt((np.power(x, 2).sum(axis=1) / number_of_points).mean())) + A = np.dot(A, np.eye(dimensionality) * scaling) + + xfrm = ants.create_ants_transform(matrix=A, translation=translation, + dimension=dimensionality, center=center_fixed) + + return xfrm + + elif transform_type == "bspline": + + bspline_displacement_field = ants.fit_bspline_displacement_field( + displacement_origins=fixed_points, + displacements=moving_points - fixed_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=enforce_stationary_boundary, + rasterize_points=rasterize_points) + + xfrm = ants.transform_from_displacement_field(bspline_displacement_field) + + return xfrm + + elif transform_type == "tps": + + tps_displacement_field = ants.fit_thin_plate_spline_displacement_field( + displacement_origins=fixed_points, + displacements=moving_points - fixed_points, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction) + + xfrm = ants.transform_from_displacement_field(tps_displacement_field) + + return xfrm + + elif transform_type == "diffeo": + + if verbose: + start_total_time = time.time() + + updated_fixed_points = np.empty_like(fixed_points) + updated_fixed_points[:] = fixed_points + + total_field = create_zero_displacement_field(domain_image) + total_field_xfrm = None + + error_values = [] + for i in range(number_of_compositions): + + if verbose: + start_time = time.time() + + update_field = ants.fit_bspline_displacement_field( + displacement_origins=updated_fixed_points, + displacements=moving_points - updated_fixed_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=True, + rasterize_points=rasterize_points + ) + + update_field = update_field * composition_step_size + if sigma > 0: + update_field = ants.smooth_image(update_field, sigma) + + total_field = ants.compose_displacement_fields(update_field, total_field) + total_field_xfrm = ants.transform_from_displacement_field(total_field) + + if i < number_of_compositions - 1: + for j in range(updated_fixed_points.shape[0]): + updated_fixed_points[j,:] = total_field_xfrm.apply_to_point(tuple(fixed_points[j,:])) + + error_values.append(np.mean(np.sqrt(np.sum(np.square(updated_fixed_points - moving_points), axis=1, keepdims=True)))) + convergence_value = convergence_monitoring(error_values) + if verbose: + end_time = time.time() + diff_time = end_time - start_time + print("Composition " + str(i) + ": error = " + str(error_values[-1]) + + " (convergence = " + str(convergence_value) + ", elapsed time = " + str(diff_time) + ")") + if not convergence_value is None and convergence_value <= convergence_threshold: + break + + if verbose: + end_total_time = time.time() + diff_total_time = end_total_time - start_total_time + print("Total elapsed time = " + str(diff_total_time) + ".") + + return(total_field_xfrm) + + elif transform_type == "syn": + + if verbose: + start_total_time = time.time() + + updated_fixed_points = np.empty_like(fixed_points) + updated_fixed_points[:] = fixed_points + updated_moving_points = np.empty_like(moving_points) + updated_moving_points[:] = moving_points + + total_field_fixed_to_middle = create_zero_displacement_field(domain_image) + total_inverse_field_fixed_to_middle = create_zero_displacement_field(domain_image) + + total_field_moving_to_middle = create_zero_displacement_field(domain_image) + total_inverse_field_moving_to_middle = create_zero_displacement_field(domain_image) + + error_values = [] + for i in range(number_of_compositions): + + if verbose: + start_time = time.time() + + update_field_fixed_to_middle = ants.fit_bspline_displacement_field( + displacement_origins=updated_fixed_points, + displacements=updated_moving_points - updated_fixed_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=True, + rasterize_points=rasterize_points + ) + + update_field_moving_to_middle = ants.fit_bspline_displacement_field( + displacement_origins=updated_moving_points, + displacements=updated_fixed_points - updated_moving_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=True, + rasterize_points=rasterize_points + ) + + update_field_fixed_to_middle = update_field_fixed_to_middle * composition_step_size + update_field_moving_to_middle = update_field_moving_to_middle * composition_step_size + if sigma > 0: + update_field_fixed_to_middle = ants.smooth_image(update_field_fixed_to_middle, sigma) + update_field_moving_to_middle = ants.smooth_image(update_field_moving_to_middle, sigma) + + # Add the update field to both forward displacement fields. + + total_field_fixed_to_middle = ants.compose_displacement_fields(update_field_fixed_to_middle, total_field_fixed_to_middle) + total_field_moving_to_middle = ants.compose_displacement_fields(update_field_moving_to_middle, total_field_moving_to_middle) + + # Iteratively estimate the inverse fields. + + total_inverse_field_fixed_to_middle = ants.invert_displacement_field(total_field_fixed_to_middle, total_inverse_field_fixed_to_middle) + total_inverse_field_moving_to_middle = ants.invert_displacement_field(total_field_moving_to_middle, total_inverse_field_moving_to_middle) + + total_field_fixed_to_middle = ants.invert_displacement_field(total_inverse_field_fixed_to_middle, total_field_fixed_to_middle) + total_field_moving_to_middle = ants.invert_displacement_field(total_inverse_field_moving_to_middle, total_field_moving_to_middle) + + total_field_fixed_to_middle_xfrm = ants.transform_from_displacement_field(total_field_fixed_to_middle) + total_field_moving_to_middle_xfrm = ants.transform_from_displacement_field(total_field_moving_to_middle) + + total_inverse_field_fixed_to_middle_xfrm = ants.transform_from_displacement_field(total_inverse_field_fixed_to_middle) + total_inverse_field_moving_to_middle_xfrm = ants.transform_from_displacement_field(total_inverse_field_moving_to_middle) + + if i < number_of_compositions - 1: + for j in range(updated_fixed_points.shape[0]): + updated_fixed_points[j,:] = total_field_fixed_to_middle_xfrm.apply_to_point(tuple(fixed_points[j,:])) + updated_moving_points[j,:] = total_field_moving_to_middle_xfrm.apply_to_point(tuple(moving_points[j,:])) + + error_values.append(np.mean(np.sqrt(np.sum(np.square(updated_fixed_points - updated_moving_points), axis=1, keepdims=True)))) + convergence_value = convergence_monitoring(error_values) + if verbose: + end_time = time.time() + diff_time = end_time - start_time + print("Composition " + str(i) + ": error = " + str(error_values[-1]) + + " (convergence = " + str(convergence_value) + ", elapsed time = " + str(diff_time) + ")") + if not convergence_value is None and convergence_value <= convergence_threshold: + break + + total_forward_field = ants.compose_displacement_fields(total_inverse_field_moving_to_middle, total_field_fixed_to_middle) + total_forward_xfrm = ants.transform_from_displacement_field(total_forward_field) + total_inverse_field = ants.compose_displacement_fields(total_inverse_field_fixed_to_middle, total_field_moving_to_middle) + total_inverse_xfrm = ants.transform_from_displacement_field(total_inverse_field) + + if verbose: + end_total_time = time.time() + diff_total_time = end_total_time - start_total_time + print("Total elapsed time = " + str(diff_total_time) + ".") + + return_dict = {'forward_transform' : total_forward_xfrm, + 'inverse_transform' : total_inverse_xfrm, + 'fixed_to_middle_transform' : total_field_fixed_to_middle_xfrm, + 'middle_to_fixed_transform' : total_inverse_field_fixed_to_middle_xfrm, + 'moving_to_middle_transform' : total_field_moving_to_middle_xfrm, + 'middle_to_moving_transform' : total_inverse_field_moving_to_middle_xfrm + } + return(return_dict) + + elif transform_type == "tv" or transform_type == "time-varying": + + if verbose: + start_total_time = time.time() + + updated_fixed_points = np.empty_like(fixed_points) + updated_fixed_points[:] = fixed_points + updated_moving_points = np.empty_like(moving_points) + updated_moving_points[:] = moving_points + + velocity_field = create_zero_velocity_field(domain_image, number_of_time_steps) + velocity_field_array = velocity_field.numpy() + + last_update_derivative_field = create_zero_velocity_field(domain_image, number_of_time_steps) + last_update_derivative_field_array = last_update_derivative_field.numpy() + + error_values = [] + for i in range(number_of_compositions): + + if verbose: + start_time = time.time() + + update_derivative_field = create_zero_velocity_field(domain_image, number_of_time_steps) + update_derivative_field_array = update_derivative_field.numpy() + + average_error = 0.0 + for n in range(number_of_time_steps): + + t = n / (number_of_time_steps - 1.0) + + if n > 0: + integrated_forward_field = ants.integrate_velocity_field(velocity_field, 0.0, t, number_of_integration_steps) + integrated_forward_field_xfrm = ants.transform_from_displacement_field(integrated_forward_field) + for j in range(updated_fixed_points.shape[0]): + updated_fixed_points[j,:] = integrated_forward_field_xfrm.apply_to_point(tuple(fixed_points[j,:])) + else: + updated_fixed_points[:] = fixed_points + + if n < number_of_time_steps - 1: + integrated_inverse_field = ants.integrate_velocity_field(velocity_field, 1.0, t, number_of_integration_steps) + integrated_inverse_field_xfrm = ants.transform_from_displacement_field(integrated_inverse_field) + for j in range(updated_moving_points.shape[0]): + updated_moving_points[j,:] = integrated_inverse_field_xfrm.apply_to_point(tuple(moving_points[j,:])) + else: + updated_moving_points[:] = moving_points + + update_derivative_field_at_timepoint = ants.fit_bspline_displacement_field( + displacement_origins=updated_fixed_points, + displacements=updated_moving_points - updated_fixed_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=True, + rasterize_points=rasterize_points + ) + + if sigma > 0: + update_derivative_field_at_timepoint = ants.smooth_image(update_derivative_field_at_timepoint, sigma) + + update_derivative_field_at_timepoint_array = update_derivative_field_at_timepoint.numpy() + grad_norms = np.sqrt(np.sum(np.square(update_derivative_field_at_timepoint_array), axis=-1, keepdims=False)) + max_norm = np.amax(grad_norms) + median_norm = np.median(grad_norms) + if verbose: + print(" integration point " + str(t) + ": max_norm = " + str(max_norm) + ", median_norm = " + str(median_norm)) + update_derivative_field_at_timepoint_array /= max_norm + if domain_image.dimension == 2: + update_derivative_field_array[:,:,n,:] = update_derivative_field_at_timepoint_array + elif domain_image.dimension == 3: + update_derivative_field_array[:,:,:,n,:] = update_derivative_field_at_timepoint_array + + rmse = np.mean(np.sqrt(np.sum(np.square(updated_moving_points - updated_fixed_points), axis=1, keepdims=True))) + average_error = (average_error * n + rmse) / (n + 1) + + update_derivative_field_array = (update_derivative_field_array + last_update_derivative_field_array) * 0.5 + last_update_derivative_field_array = np.empty_like(update_derivative_field_array) + last_update_derivative_field_array[:] = update_derivative_field_array + + velocity_field_array = velocity_field_array + update_derivative_field_array * composition_step_size + velocity_field = ants.from_numpy(velocity_field_array, origin=velocity_field.origin, + spacing=velocity_field.spacing, direction=velocity_field.direction, + has_components=True) + + error_values.append(average_error) + convergence_value = convergence_monitoring(error_values) + if verbose: + end_time = time.time() + diff_time = end_time - start_time + print("Composition " + str(i) + ": error = " + str(error_values[-1]) + + " (convergence = " + str(convergence_value) + ", elapsed time = " + str(diff_time) + ")") + if not convergence_value is None and convergence_value <= convergence_threshold: + break + + forward_xfrm = ants.transform_from_displacement_field(ants.integrate_velocity_field(velocity_field, 0.0, 1.0, number_of_integration_steps)) + inverse_xfrm = ants.transform_from_displacement_field(ants.integrate_velocity_field(velocity_field, 1.0, 0.0, number_of_integration_steps)) + + if verbose: + end_total_time = time.time() + diff_total_time = end_total_time - start_total_time + print("Total elapsed time = " + str(diff_total_time) + ".") + + return_dict = {'forward_transform': forward_xfrm, + 'inverse_transform': inverse_xfrm, + 'velocity_field': velocity_field} + return(return_dict) + + else: + raise ValueError("Unrecognized transform_type.") + + +def fit_time_varying_transform_to_point_sets(point_sets, + time_points=None, + initial_velocity_field=None, + number_of_time_steps=None, + domain_image=None, + number_of_fitting_levels=4, + mesh_size=1, + spline_order=3, + displacement_weights=None, + number_of_compositions=10, + composition_step_size=0.5, + number_of_integration_steps=100, + sigma=0.0, + convergence_threshold=1e-6, + rasterize_points=False, + verbose=False + ): + """ + + Estimate a time-varying transform from corresponding point sets (> 2). + + ANTsR function: fitTimeVaryingTransformToPointSets + + Arguments + --------- + point_sets : list of arrays + Corresponding points across sets specified in physical space as a n x d matrix where n + is the number of points and d is the dimensionality. + + time_points : array of ordered scalars between 0 and 1 + Set of scalar values, one for each point-set, designating its time position in the velocity + flow. If not set, it defaults to equal spacing between 0 and 1. + + initial_velocity_field : initial ANTs velocity field + Optional velocity field for initializing optimization. Overrides the number of integration + points. + + number_of_time_steps : integer + Time-varying velocity field parameter. Needs to be equal to or greater than the number of + point sets. If not specified, it defaults to the number of point sets. + + domain_image : ANTs image + Defines physical domain of the nonlinear transform. Must be defined. + + number_of_fitting_levels : integer + Integer specifying the number of fitting levels for the B-spline interpolation of the + displacement field. + + mesh_size : integer or array + Defines the mesh size at the initial fitting level for the B-spline interpolation of the + displacement field.. + + spline_order : integer + Spline order of the B-spline displacement field. + + displacement_weights : array + Defines the individual weighting of the corresponding scattered data value. Default = NULL + meaning all displacements are weighted the same. + + number_of_compositions : integer + Total number of compositions. + + composition_step_size : scalar + Scalar multiplication factor of the weighting of the update field. + + number_of_integration_steps : scalar + Number of steps used for integrating the velocity field. + + sigma : scalar + Gaussian smoothing standard deviation of the update field (in mm). + + convergence_threshold : scalar + Composition-based convergence parameter using a window size of 10 values. + + rasterize_points : boolean + Use nearest neighbor rasterization of points for estimating the update field (potential + speed-up). Default = False. + + verbose : bool + Print progress to the screen. + + Returns + ------- + + ANTs transform + + Example + ------- + >>> import ants + >>> import numpy as np + """ + + def create_zero_velocity_field(domain_image, number_of_time_points=2): + field_array = np.zeros((*domain_image.shape, number_of_time_points, domain_image.dimension)) + origin = (*domain_image.origin, 0.0) + spacing = (*domain_image.spacing, 1.0) + direction = np.eye(domain_image.dimension + 1) + direction[0:domain_image.dimension,0:domain_image.dimension] = domain_image.direction + field = ants.from_numpy(field_array, origin=origin, spacing=spacing, direction=direction, + has_components=True) + return(field) + + if not isinstance(point_sets, list): + raise ValueError("point_sets should be a list of corresponding point sets.") + + number_of_point_sets = len(point_sets) + + if time_points is not None and len(time_points) != number_of_point_sets: + raise ValueError("The number of time points should be the same as the number of point sets.") + + if time_points is None: + time_points = np.linspace(0.0, 1.0, number_of_point_sets) + time_points = np.array(time_points) + + if np.any(time_points < 0.0) or np.any(time_points > 1.0): + raise ValueError("time point values should be between 0 and 1.") + + if number_of_point_sets < 3: + raise ValueError("Expecting three or greater point sets.") + + if domain_image is None: + raise ValueError("Domain image needs to be specified.") + + number_of_points = point_sets[0].shape[0] + dimensionality = point_sets[0].shape[1] + for i in range(1, number_of_point_sets): + if point_sets[i].shape[0] != number_of_points: + raise ValueError("Point sets should match in terms of the number of points.") + if point_sets[i].shape[1] != dimensionality: + raise ValueError("Point sets should match in terms of dimensionality.") + + if verbose: + start_total_time = time.time() + + updated_fixed_points = np.zeros(point_sets[0].shape) + updated_moving_points = np.zeros(point_sets[0].shape) + + velocity_field = None + if initial_velocity_field is None: + if number_of_time_steps is None: + number_of_time_steps = len(time_points) + if number_of_time_steps < number_of_point_sets: + raise ValueError("The number of integration points should be at least as great as the number of point sets.") + velocity_field = create_zero_velocity_field(domain_image, number_of_time_steps) + else: + velocity_field = ants.image_clone(initial_velocity_field) + number_of_time_steps = initial_velocity_field.shape[-1] + velocity_field_array = velocity_field.numpy() + + last_update_derivative_field = create_zero_velocity_field(domain_image, number_of_time_steps) + last_update_derivative_field_array = last_update_derivative_field.numpy() + + error_values = [] + for i in range(number_of_compositions): + + if verbose: + start_time = time.time() + + update_derivative_field = create_zero_velocity_field(domain_image, number_of_time_steps) + update_derivative_field_array = update_derivative_field.numpy() + + average_error = 0.0 + for n in range(number_of_time_steps): + + t = n / (number_of_time_steps - 1.0) + + t_index = 0 + for j in range(1, number_of_point_sets): + if time_points[j-1] <= t and time_points[j] >= t: + t_index = j + break + + if n > 0 and n < number_of_time_steps - 1 and time_points[t_index-1] == t: + updated_fixed_points[:] = point_sets[t_index-1] + integrated_inverse_field = ants.integrate_velocity_field(velocity_field, time_points[t_index], t, number_of_integration_steps) + integrated_inverse_field_xfrm = ants.transform_from_displacement_field(integrated_inverse_field) + for j in range(updated_moving_points.shape[0]): + updated_moving_points[j,:] = integrated_inverse_field_xfrm.apply_to_point(tuple(point_sets[t_index][j,:])) + + update_derivative_field_at_timepoint_forward = ants.fit_bspline_displacement_field( + displacement_origins=updated_fixed_points, + displacements=updated_moving_points - updated_fixed_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=True, + rasterize_points=rasterize_points + ) + + updated_moving_points[:] = point_sets[t_index-1] + integrated_forward_field = ants.integrate_velocity_field(velocity_field, time_points[t_index-2], t, number_of_integration_steps) + integrated_forward_field_xfrm = ants.transform_from_displacement_field(integrated_forward_field) + for j in range(updated_fixed_points.shape[0]): + updated_fixed_points[j,:] = integrated_forward_field_xfrm.apply_to_point(tuple(point_sets[t_index-2][j,:])) + + update_derivative_field_at_timepoint_back = ants.fit_bspline_displacement_field( + displacement_origins=updated_fixed_points, + displacements=updated_moving_points - updated_fixed_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=True, + rasterize_points=rasterize_points + ) + + update_derivative_field_at_timepoint = (update_derivative_field_at_timepoint_forward + + update_derivative_field_at_timepoint_back) / 2.0 + + else: + if t == 0.0 and time_points[t_index-1] == 0.0: + updated_fixed_points[:] = point_sets[0] + else: + integrated_forward_field = ants.integrate_velocity_field(velocity_field, time_points[t_index-1], t, number_of_integration_steps) + integrated_forward_field_xfrm = ants.transform_from_displacement_field(integrated_forward_field) + for j in range(updated_fixed_points.shape[0]): + updated_fixed_points[j,:] = integrated_forward_field_xfrm.apply_to_point(tuple(point_sets[t_index-1][j,:])) + + if t == 1.0 and time_points[t_index] == 1.0: + updated_moving_points[:] = point_sets[-1] + else: + integrated_inverse_field = ants.integrate_velocity_field(velocity_field, time_points[t_index], t, number_of_integration_steps) + integrated_inverse_field_xfrm = ants.transform_from_displacement_field(integrated_inverse_field) + for j in range(updated_moving_points.shape[0]): + updated_moving_points[j,:] = integrated_inverse_field_xfrm.apply_to_point(tuple(point_sets[t_index][j,:])) + + update_derivative_field_at_timepoint = ants.fit_bspline_displacement_field( + displacement_origins=updated_fixed_points, + displacements=updated_moving_points - updated_fixed_points, + displacement_weights=displacement_weights, + origin=domain_image.origin, + spacing=domain_image.spacing, + size=domain_image.shape, + direction=domain_image.direction, + number_of_fitting_levels=number_of_fitting_levels, + mesh_size=mesh_size, + spline_order=spline_order, + enforce_stationary_boundary=True, + rasterize_points=rasterize_points + ) + + if sigma > 0: + update_derivative_field_at_timepoint = ants.smooth_image(update_derivative_field_at_timepoint, sigma) + + update_derivative_field_at_timepoint_array = update_derivative_field_at_timepoint.numpy() + grad_norms = np.sqrt(np.sum(np.square(update_derivative_field_at_timepoint_array), axis=-1, keepdims=False)) + max_norm = np.amax(grad_norms) + median_norm = np.median(grad_norms) + if verbose: + print(" integration point " + str(t) + ": max_norm = " + str(max_norm) + ", median_norm = " + str(median_norm)) + update_derivative_field_at_timepoint_array /= max_norm + if domain_image.dimension == 2: + update_derivative_field_array[:,:,n,:] = update_derivative_field_at_timepoint_array + elif domain_image.dimension == 3: + update_derivative_field_array[:,:,:,n,:] = update_derivative_field_at_timepoint_array + + rmse = np.mean(np.sqrt(np.sum(np.square(updated_moving_points - updated_fixed_points), axis=1, keepdims=True))) + average_error = (average_error * n + rmse) / (n + 1) + + update_derivative_field_array = (update_derivative_field_array + last_update_derivative_field_array) * 0.5 + last_update_derivative_field_array = np.empty_like(update_derivative_field_array) + last_update_derivative_field_array[:] = update_derivative_field_array + + velocity_field_array += (update_derivative_field_array * composition_step_size) + velocity_field = ants.from_numpy(velocity_field_array, origin=velocity_field.origin, + spacing=velocity_field.spacing, direction=velocity_field.direction, + has_components=True) + + error_values.append(average_error) + convergence_value = convergence_monitoring(error_values) + if verbose: + end_time = time.time() + diff_time = end_time - start_time + print("Composition " + str(i) + ": error = " + str(error_values[-1]) + + " (convergence = " + str(convergence_value) + ", elapsed time = " + str(diff_time) + ")") + if not convergence_value is None and convergence_value <= convergence_threshold: + break + + forward_xfrm = ants.transform_from_displacement_field(ants.integrate_velocity_field(velocity_field, 0.0, 1.0, number_of_integration_steps)) + inverse_xfrm = ants.transform_from_displacement_field(ants.integrate_velocity_field(velocity_field, 1.0, 0.0, number_of_integration_steps)) + + if verbose: + end_total_time = time.time() + diff_total_time = end_total_time - start_total_time + print("Total elapsed time = " + str(diff_total_time) + ".") + + return_dict = {'forward_transform': forward_xfrm, + 'inverse_transform': inverse_xfrm, + 'velocity_field': velocity_field} + return(return_dict) + + + + diff --git a/MindEyeV2/antspy/ants/registration/registration.py b/MindEyeV2/antspy/ants/registration/registration.py new file mode 100644 index 0000000000000000000000000000000000000000..2c6739a19e77f32a17938021f87e63dc55123df8 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/registration.py @@ -0,0 +1,1953 @@ +""" +ANTsPy Registration +""" +__all__ = ["registration", + "motion_correction", + "label_image_registration"] + +import numpy as np +from tempfile import mktemp +import glob +import re +import pandas as pd +import itertools + +import ants +from ants.internal import get_lib_fn, get_pointer_string, process_arguments + +def registration( + fixed, + moving, + type_of_transform="SyN", + initial_transform=None, + outprefix="", + mask=None, + moving_mask=None, + mask_all_stages=False, + grad_step=0.2, + flow_sigma=3, + total_sigma=0, + aff_metric="mattes", + aff_sampling=32, + aff_random_sampling_rate=0.2, + syn_metric="mattes", + syn_sampling=32, + reg_iterations=(40, 20, 0), + aff_iterations=(2100, 1200, 1200, 10), + aff_shrink_factors=(6, 4, 2, 1), + aff_smoothing_sigmas=(3, 2, 1, 0), + write_composite_transform=False, + random_seed=None, + verbose=False, + multivariate_extras=None, + restrict_transformation=None, + smoothing_in_mm=False, + singleprecision=True, + **kwargs +): + """ + Register a pair of images either through the full or simplified + interface to the ANTs registration method. + + ANTsR function: `antsRegistration` + + Arguments + --------- + fixed : ANTsImage + fixed image to which we register the moving image. + + moving : ANTsImage + moving image to be mapped to fixed space. + + type_of_transform : string + A linear or non-linear registration type. Mutual information metric by default. + See Notes below for more. + + initial_transform : list of strings (optional) + transforms to prepend. If None, a translation is computed to align the image centers of mass. + To use an identity transform, set this to 'Identity'. + + outprefix : string + output will be named with this prefix. + + mask : ANTsImage (optional) + Registration metric mask in the fixed image space. + + moving_mask : ANTsImage (optional) + Registration metric mask in the moving image space. + + mask_all_stages : boolean + If true, apply metric mask(s) to all registration stages, instead of just the final stage. + + grad_step : scalar + gradient step size (not for all tx) + + flow_sigma : scalar + smoothing for update field + At each iteration, the similarity metric and gradient is calculated. + That gradient field is also called the update field and is smoothed + before composing with the total field (i.e., the estimate of the total + transform at that iteration). This total field can also be smoothed + after each iteration. + + total_sigma : scalar + smoothing for total field + + aff_metric : string + the metric for the affine part (GC, mattes, meansquares) + + aff_sampling : scalar + number of bins for the mutual information metric + + aff_random_sampling_rate : scalar + the fraction of points used to estimate the metric. this can impact + speed but also reproducibility and/or accuracy. + + syn_metric : string + the metric for the syn part (CC, mattes, meansquares, demons) + + syn_sampling : scalar + the nbins or radius parameter for the syn metric + + reg_iterations : list/tuple of integers + vector of iterations for syn. we will set the smoothing and multi-resolution parameters based on the length of this vector. + + aff_iterations : list/tuple of integers + vector of iterations for low-dimensional (translation, rigid, affine) registration. + + aff_shrink_factors : list/tuple of integers + vector of multi-resolution shrink factors for low-dimensional (translation, rigid, affine) registration. + + aff_smoothing_sigmas : list/tuple of integers + vector of multi-resolution smoothing factors for low-dimensional (translation, rigid, affine) registration. + + random_seed : integer + random seed to improve reproducibility. note that the number of ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS should be 1 if you want perfect reproducibility. + + write_composite_transform : boolean + Boolean specifying whether or not the composite transform (and its inverse, if it exists) should be written to an hdf5 composite file. This is false by default so that only the transform for each stage is written to file. + + verbose : boolean + request verbose output (useful for debugging) + + multivariate_extras : additional metrics for multi-metric registration + list of additional images and metrics which will + trigger the use of multiple metrics in the registration + process in the deformable stage. Each multivariate metric needs 5 + entries: name of metric, fixed, moving, weight, + samplingParam. the list of lists should be of the form ( ( + "nameOfMetric2", img, img, weight, metricParam ) ). Another + example would be ( ( "MeanSquares", f2, m2, 0.5, 0 + ), ( "CC", f2, m2, 0.5, 2 ) ) . This is only compatible + with the SyNOnly or antsRegistrationSyN* transformations. + + restrict_transformation : This option allows the user to restrict the + optimization of the displacement field, translation, rigid or + affine transform on a per-component basis. For example, if + one wants to limit the deformation or rotation of 3-D volume + to the first two dimensions, this is possible by specifying a + weight vector of ‘(1,1,0)’ for a 3D deformation field or + ‘(1,1,0,1,1,0)’ for a rigid transformation. Restriction + currently only works if there are no preceding + transformations. + + smoothing_in_mm : boolean ; currently only impacts low dimensional registration + + singleprecision : boolean + if True, use float32 for computations. This is useful for reducing memory + usage for large datasets, at the cost of precision. + + kwargs : keyword args + extra arguments + + Returns + ------- + dict containing follow key/value pairs: + `warpedmovout`: Moving image warped to space of fixed image. + `warpedfixout`: Fixed image warped to space of moving image. + `fwdtransforms`: Transforms to move from moving to fixed image. + `invtransforms`: Transforms to move from fixed to moving image. + + Notes + ----- + type_of_transform can be one of: + - "Translation": Translation transformation. + - "Rigid": Rigid transformation: Only rotation and translation. + - "Similarity": Similarity transformation: scaling, rotation and translation. + - "QuickRigid": Rigid transformation: Only rotation and translation. + May be useful for quick visualization fixes.' + - "DenseRigid": Rigid transformation: Only rotation and translation. + Employs dense sampling during metric estimation.' + - "BOLDRigid": Rigid transformation: Parameters typical for BOLD to + BOLD intrasubject registration'.' + - "Affine": Affine transformation: Rigid + scaling. + - "AffineFast": Fast version of Affine. + - "BOLDAffine": Affine transformation: Parameters typical for BOLD to + BOLD intrasubject registration'.' + - "TRSAA": translation, rigid, similarity, affine (twice). please set + regIterations if using this option. this would be used in + cases where you want a really high quality affine mapping + (perhaps with mask). + - "Elastic": Elastic deformation: Affine + deformable. + - "ElasticSyN": Symmetric normalization: Affine + deformable + transformation, with mutual information as optimization + metric and elastic regularization. + - "SyN": Symmetric normalization: Affine + deformable transformation, + with mutual information as optimization metric. + - "SyNRA": Symmetric normalization: Rigid + Affine + deformable + transformation, with mutual information as optimization metric. + - "SyNOnly": Symmetric normalization with no rigid or affine stages. + Uses mutual information as optimization metric. Affine alignment is + from the initial_transform arg, either provide the .mat from linear + registration or use initial_transform='Identity' if the images are + already affinely aligned. + Can be useful if you want to run an unmasked affine followed by + masked deformable registration. + - "SyNCC": SyN, but with cross-correlation as the metric. + - "SyNabp": SyN optimized for abpBrainExtraction. + - "SyNBold": SyN, but optimized for registrations between BOLD and T1 images. + - "SyNBoldAff": SyN, but optimized for registrations between BOLD + and T1 images, with additional affine step. + - "SyNAggro": SyN, but with more aggressive registration + (fine-scale matching and more deformation). + Takes more time than SyN. + - "TV[n]": time-varying diffeomorphism with where 'n' indicates number of + time points in velocity field discretization. The initial transform + should be computed, if needed, in a separate call to ants.registration. + - "TVMSQ": time-varying diffeomorphism with mean square metric + - "TVMSQC": time-varying diffeomorphism with mean square metric for very large deformation + - "antsRegistrationSyN[x]": recreation of the antsRegistrationSyN.sh script in ANTs + where 'x' is one of the transforms available: + t: translation (1 stage) + r: rigid (1 stage) + a: rigid + affine (2 stages) + s: rigid + affine + deformable syn (3 stages) + sr: rigid + deformable syn (2 stages) + so: deformable syn only (1 stage) + b: rigid + affine + deformable b-spline syn (3 stages) + br: rigid + deformable b-spline syn (2 stages) + bo: deformable b-spline syn only (1 stage) + - "antsRegistrationSyNQuick[x]": recreation of the antsRegistrationSyNQuick.sh script in ANTs. + x options as above. + - "antsRegistrationSyNRepro[x]": reproducible registration. x options as above. + - "antsRegistrationSyNQuickRepro[x]": quick reproducible registration. x options as above. + + Example + ------- + >>> import ants + >>> fi = ants.image_read(ants.get_ants_data('r16')) + >>> mi = ants.image_read(ants.get_ants_data('r64')) + >>> fi = ants.resample_image(fi, (60,60), 1, 0) + >>> mi = ants.resample_image(mi, (60,60), 1, 0) + >>> mytx = ants.registration(fixed=fi, moving=mi, type_of_transform = 'SyN' ) + >>> mytx = ants.registration(fixed=fi, moving=mi, type_of_transform = 'antsRegistrationSyN[t]' ) + >>> mytx = ants.registration(fixed=fi, moving=mi, type_of_transform = 'antsRegistrationSyN[b]' ) + >>> mytx = ants.registration(fixed=fi, moving=mi, type_of_transform = 'antsRegistrationSyN[s]' ) + """ + if isinstance(fixed, list) and (moving is None): + processed_args = process_arguments(fixed) + libfn = get_lib_fn("antsRegistration") + reg_exit = libfn(processed_args) + if (reg_exit != 0): + raise RuntimeError(f"Registration failed with error code {reg_exit}") + else: + return 0 + + if not (ants.is_image(fixed) and ants.is_image(moving)): + raise ValueError("Fixed and moving images must be ANTsImage objects") + + if type_of_transform == "": + type_of_transform = "SyN" + + if isinstance(type_of_transform, (tuple, list)) and (len(type_of_transform) == 1): + type_of_transform = type_of_transform[0] + + if (outprefix == "") or len(outprefix) == 0: + outprefix = mktemp() + + if np.sum(np.isnan(fixed.numpy())) > 0: + raise ValueError("fixed image has NaNs - replace these") + if np.sum(np.isnan(moving.numpy())) > 0: + raise ValueError("moving image has NaNs - replace these") + + if fixed.dimension != moving.dimension: + raise ValueError("Fixed and moving image dimensions are not the same.") + # ---------------------------- + + myiterations = aff_iterations + args = [fixed, moving, type_of_transform, outprefix] + myf_aff = "6x4x2x1" # old fixed params + mys_aff = "3x2x1x0" # old fixed params + if ( + type(aff_shrink_factors) is int + or type(aff_smoothing_sigmas) is int + or type(aff_iterations) is int + ): + if type(aff_smoothing_sigmas) is not int: + raise ValueError("aff_smoothing_sigmas should be a single integer.") + if type(aff_iterations) is not int: + raise ValueError("aff_iterations should be a single integer.") + if type(aff_shrink_factors) is not int: + raise ValueError("aff_shrink_factors should be a single integer.") + myf_aff = aff_shrink_factors + mys_aff = aff_smoothing_sigmas + myiterations = aff_iterations + + if restrict_transformation is not None: + if type(restrict_transformation) is tuple: + restrict_transformationchar = "x".join([str(ri) for ri in restrict_transformation]) + + if type(aff_shrink_factors) is tuple: + myf_aff = "x".join([str(ri) for ri in aff_shrink_factors]) + mys_aff = "x".join([str(ri) for ri in aff_smoothing_sigmas]) + myiterations = "x".join([str(ri) for ri in aff_iterations]) + if len(aff_iterations) != len(aff_smoothing_sigmas): + raise ValueError( + "aff_iterations length should equal aff_smoothing_sigmas length." + ) + if len(aff_iterations) != len(aff_shrink_factors): + raise ValueError( + "aff_iterations length should equal aff_shrink_factors length." + ) + if len(aff_shrink_factors) != len(aff_smoothing_sigmas): + raise ValueError( + "aff_shrink_factors length should equal aff_smoothing_sigmas length." + ) + + if type_of_transform == "AffineFast": + type_of_transform = "Affine" + myiterations = "2100x1200x0x0" + if type_of_transform == "BOLDAffine": + type_of_transform = "Affine" + myf_aff = "2x1" + mys_aff = "1x0" + myiterations = "100x20" + if type_of_transform == "QuickRigid": + type_of_transform = "Rigid" + myiterations = "20x20x0x0" + if type_of_transform == "DenseRigid": + type_of_transform = "Rigid" + aff_random_sampling_rate = 1.0 + if type_of_transform == "BOLDRigid": + type_of_transform = "Rigid" + myf_aff = "2x1" + mys_aff = "1x0" + myiterations = "100x20" + + if smoothing_in_mm: + mys_aff = mys_aff + 'mm' + + mysyn = "SyN[%f,%f,%f]" % (grad_step, flow_sigma, total_sigma) + if type_of_transform == "Elastic": + mysyn = "GaussianDisplacementField[%f,%f,%f]" % (grad_step, flow_sigma, total_sigma) + itlen = len(reg_iterations) # NEED TO CHECK THIS + if itlen == 0: + smoothingsigmas = 0 + shrinkfactors = 1 + synits = reg_iterations + else: + smoothingsigmas = np.arange(0, itlen)[::-1].astype( + "float32" + ) # NEED TO CHECK THIS + shrinkfactors = 2 ** smoothingsigmas + shrinkfactors = shrinkfactors.astype("int") + smoothingsigmas = "x".join([str(ss)[0] for ss in smoothingsigmas]) + shrinkfactors = "x".join([str(ss) for ss in shrinkfactors]) + synits = "x".join([str(ri) for ri in reg_iterations]) + + inpixeltype = fixed.pixeltype + output_pixel_type = 'float' if singleprecision else 'double' + + tvTypes = [ + "TV[1]", + "TV[2]", + "TV[3]", + "TV[4]", + "TV[5]", + "TV[6]", + "TV[7]", + "TV[8]", + ] + allowable_tx = { + "SyNBold", + "SyNBoldAff", + "ElasticSyN", + "Elastic", + "SyN", + "SyNRA", + "SyNOnly", + "SyNAggro", + "SyNCC", + "TRSAA", + "SyNabp", + "SyNLessAggro", + "TV[1]", + "TV[2]", + "TV[3]", + "TV[4]", + "TV[5]", + "TV[6]", + "TV[7]", + "TV[8]", + "TVMSQ", + "TVMSQC", + "Rigid", + "Similarity", + "Translation", + "Affine", + "AffineFast", + "BOLDAffine", + "QuickRigid", + "DenseRigid", + "BOLDRigid" + } + ttexists = type_of_transform in allowable_tx + + # Perform checking of antsRegistrationSyN transforms later + if not "antsRegistrationSyN" in type_of_transform and not ttexists: + raise ValueError(f'{type_of_transform} does not exist') + + initx = initial_transform + if isinstance(initx, str): + initx = [initx] + # if isinstance(initx, ANTsTransform): + # tempTXfilename = tempfile( fileext = '.mat' ) + # initx = invertAntsrTransform( initialTransform ) + # initx = invertAntsrTransform( initx ) + # writeAntsrTransform( initx, tempTXfilename ) + # initx = tempTXfilename + moving = moving.clone(output_pixel_type) + fixed = fixed.clone(output_pixel_type) + # NOTE: this may be better for general purpose applications: TBD +# moving = ants.iMath( moving.clone("float"), "Normalize" ) +# fixed = ants.iMath( fixed.clone("float"), "Normalize" ) + warpedfixout = moving.clone() + warpedmovout = fixed.clone() + f = get_pointer_string(fixed) + m = get_pointer_string(moving) + wfo = get_pointer_string(warpedfixout) + wmo = get_pointer_string(warpedmovout) + if mask is not None: + mask_binary = mask != 0 + f_mask_str = get_pointer_string(mask_binary) + else: + f_mask_str = "NA" + + if moving_mask is not None: + moving_mask_binary = moving_mask != 0 + m_mask_str = get_pointer_string(moving_mask_binary) + else: + m_mask_str = "NA" + + maskopt = "[%s,%s]" % (f_mask_str, m_mask_str) + + if mask_all_stages: + earlymaskopt = maskopt; + else: + earlymaskopt = "[NA,NA]" + + if initx is None: + initx = ["[%s,%s,1]" % (f, m)] + # ------------------------------------------------------------ + if type_of_transform == "SyNBold": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Rigid[0.25]", + "-c", + "[1200x1200x100,1e-6,5]", + "-s", + "2x1x0", + "-f", + "4x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "SyNBoldAff": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Rigid[0.25]", + "-c", + "[1200x1200x100,1e-6,5]", + "-s", + "2x1x0", + "-f", + "4x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[0.25]", + "-c", + "[200x20,1e-6,5]", + "-s", + "1x0", + "-f", + "2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % (synits), + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "ElasticSyN": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[0.25]", + "-c", + "2100x1200x200x0", + "-s", + "3x2x1x0", + "-f", + "4x2x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % (synits), + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "SyN" or type_of_transform == "Elastic": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[0.25]", + "-c", + "2100x1200x1200x0", + "-s", + "3x2x1x0", + "-f", + "4x2x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "SyNRA": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Rigid[0.25]", + "-c", + "2100x1200x1200x0", + "-s", + "3x2x1x0", + "-f", + "4x2x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[0.25]", + "-c", + "2100x1200x1200x0", + "-s", + "3x2x1x0", + "-f", + "4x2x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "SyNOnly": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + ] + if multivariate_extras is not None: + metrics = [] + for kk in range(len(multivariate_extras)): + metrics.append("-m") + metricname = multivariate_extras[kk][0] + metricfixed = get_pointer_string( + multivariate_extras[kk][1] + ) + metricmov = get_pointer_string( + multivariate_extras[kk][2] + ) + metricWeight = multivariate_extras[kk][3] + metricSampling = multivariate_extras[kk][4] + metricString = "%s[%s,%s,%s,%s]" % ( + metricname, + metricfixed, + metricmov, + metricWeight, + metricSampling, + ) + metrics.append(metricString) + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + ] + args1 = [ + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + ] + for kk in range(len(metrics)): + args.append(metrics[kk]) + for kk in range(len(args1)): + args.append(args1[kk]) + args.append("-x") + args.append(maskopt) + # ------------------------------------------------------------ + elif type_of_transform == "SyNAggro": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[0.25]", + "-c", + "2100x1200x1200x100", + "-s", + "3x2x1x0", + "-f", + "4x2x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "SyNCC": + syn_metric = "CC" + syn_sampling = 4 + synits = "2100x1200x1200x20" + smoothingsigmas = "3x2x1x0" + shrinkfactors = "4x3x2x1" + mysyn = "SyN[0.15,3,0]" + + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Rigid[1]", + "-c", + "2100x1200x1200x0", + "-s", + "3x2x1x0", + "-f", + "4x4x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[1]", + "-c", + "1200x1200x100", + "-s", + "2x1x0", + "-f", + "4x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "TRSAA": + itlen = len(reg_iterations) + itlenlow = round(itlen / 2 + 0.0001) + dlen = itlen - itlenlow + _myconvlow = [2000] * itlenlow + [0] * dlen + myconvlow = "x".join([str(mc) for mc in _myconvlow]) + myconvhi = "x".join([str(r) for r in reg_iterations]) + myconvhi = "[%s,1.e-7,10]" % myconvhi + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Translation[1]", + "-c", + myconvlow, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Rigid[1]", + "-c", + myconvlow, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Similarity[1]", + "-c", + myconvlow, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[1]", + "-c", + myconvhi, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[1]", + "-c", + myconvhi, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------s + elif type_of_transform == "SyNabp": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "mattes[%s,%s,1,32,regular,0.25]" % (f, m), + "-t", + "Rigid[0.1]", + "-c", + "1000x500x250x100", + "-s", + "4x2x1x0", + "-f", + "8x4x2x1", + "-x", + earlymaskopt, + "-m", + "mattes[%s,%s,1,32,regular,0.25]" % (f, m), + "-t", + "Affine[0.1]", + "-c", + "1000x500x250x100", + "-s", + "4x2x1x0", + "-f", + "8x4x2x1", + "-x", + earlymaskopt, + "-m", + "CC[%s,%s,0.5,4]" % (f, m), + "-t", + "SyN[0.1,3,0]", + "-c", + "50x10x0", + "-s", + "2x1x0", + "-f", + "4x2x1", + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "SyNLessAggro": + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "Affine[0.25]", + "-c", + "2100x1200x1200x100", + "-s", + "3x2x1x0", + "-f", + "4x2x2x1", + "-x", + earlymaskopt, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + mysyn, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform in tvTypes: + if grad_step is None: + grad_step = 1.0 + nTimePoints = type_of_transform.split("[")[1].split("]")[0] + tvtx = ( + "TimeVaryingVelocityField[" + + str(grad_step) + + "," + + nTimePoints + + "," + + str(flow_sigma) + + ",0.0," + + str(total_sigma) + + ",0]" + ) + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + tvtx, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "0", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + elif type_of_transform == "TVMSQ": + if grad_step is None: + grad_step = 1.0 + + tvtx = "TimeVaryingVelocityField[%s, 4, 0.0,0.0, 0.5,0 ]" % str( + grad_step + ) + args = [ + "-d", + str(fixed.dimension), + # '-r', initx, + "-m", + "%s[%s,%s,1,%s]" % (syn_metric, f, m, syn_sampling), + "-t", + tvtx, + "-c", + "[%s,1e-7,8]" % synits, + "-s", + smoothingsigmas, + "-f", + shrinkfactors, + "-u", + "1", + "-z", + "0", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif type_of_transform == "TVMSQC": + if grad_step is None: + grad_step = 2.0 + + tvtx = "TimeVaryingVelocityField[%s, 8, 1.0,0.0, 0.05,0 ]" % str( + grad_step + ) + args = [ + "-d", + str(fixed.dimension), + # '-r', initx, + "-m", + "demons[%s,%s,0.5,0]" % (f, m), + "-m", + "meansquares[%s,%s,1,0]" % (f, m), + "-t", + tvtx, + "-c", + "[1200x1200x100x20x0,0,5]", + "-s", + "8x6x4x2x1vox", + "-f", + "8x6x4x2x1", + "-u", + "1", + "-z", + "0", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif ( + (type_of_transform == "Rigid") + or (type_of_transform == "Similarity") + or (type_of_transform == "Translation") + or (type_of_transform == "Affine") + ): + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-m", + "%s[%s,%s,1,%s,regular,%s]" + % (aff_metric, f, m, aff_sampling, aff_random_sampling_rate), + "-t", + "%s[0.25]" % type_of_transform, + "-c", + myiterations, + "-s", + mys_aff, + "-f", + myf_aff, + "-u", + "1", + "-z", + "1", + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + "-x", + maskopt + ] + # ------------------------------------------------------------ + elif "antsRegistrationSyN" in type_of_transform: + + do_quick = False + if "Quick" in type_of_transform: + do_quick = True + + subtype_of_transform = "s" + spline_distance = 26 + metric_parameter = 4 + if do_quick: + metric_parameter = 32 + + if "[" in type_of_transform and "]" in type_of_transform: + subtype_of_transform = type_of_transform.split("[")[1].split( + "]" + )[0] + if "," in subtype_of_transform: + subtype_of_transform_args = subtype_of_transform.split(",") + subtype_of_transform = subtype_of_transform_args[0] + if not ( subtype_of_transform == "b" + or subtype_of_transform == "br" + or subtype_of_transform == "bo" + or subtype_of_transform == "s" + or subtype_of_transform == "sr" + or subtype_of_transform == "so" ): + raise ValueError("Extra parameters are only valid for 's' or 'b' SyN transforms.") + metric_parameter = subtype_of_transform_args[1] + if len(subtype_of_transform_args) > 2: + spline_distance = subtype_of_transform_args[2] + + do_repro = False + if "Repro" in type_of_transform: + do_repro = True + + if do_quick == True: + rigid_convergence = "[1000x500x250x0,1e-6,10]" + else: + rigid_convergence = "[1000x500x250x100,1e-6,10]" + rigid_shrink_factors = "8x4x2x1" + rigid_smoothing_sigmas = "3x2x1x0vox" + + if do_quick == True: + affine_convergence = "[1000x500x250x0,1e-6,10]" + else: + affine_convergence = "[1000x500x250x100,1e-6,10]" + affine_shrink_factors = "8x4x2x1" + affine_smoothing_sigmas = "3x2x1x0vox" + + linear_metric="MI[%s,%s,1,32,Regular,0.25]" + if do_repro == True: + linear_metric="GC[%s,%s,1,1,Regular,0.25]" + + if do_quick == True: + syn_convergence = "[100x70x50x0,1e-6,10]" + metric_parameter = 32 + syn_metric = "MI[%s,%s,1,%s]" % (f, m, metric_parameter) + else: + metric_parameter = 2 + syn_convergence = "[100x70x50x20,1e-6,10]" + syn_metric = "CC[%s,%s,1,%s]" % (f, m, metric_parameter) + syn_shrink_factors = "8x4x2x1" + syn_smoothing_sigmas = "3x2x1x0vox" + + if do_quick == True and do_repro == True: + syn_convergence = "[100x70x50x0,1e-6,10]" + metric_parameter = 2 + syn_metric = "CC[%s,%s,1,%s]" % (f, m, metric_parameter) + + if random_seed is None and do_repro == True: + random_seed = str( 1 ) + + tx = "Rigid" + if subtype_of_transform == "t": + tx = "Translation" + + rigid_stage = [ + "--transform", + tx + "[0.1]", + "--metric", + linear_metric % (f, m), + "--convergence", + rigid_convergence, + "--shrink-factors", + rigid_shrink_factors, + "--smoothing-sigmas", + rigid_smoothing_sigmas, + ] + + affine_stage = [ + "--transform", + "Affine[0.1]", + "--metric", + linear_metric % (f, m), + "--convergence", + affine_convergence, + "--shrink-factors", + affine_shrink_factors, + "--smoothing-sigmas", + affine_smoothing_sigmas, + ] + + if subtype_of_transform == "sr" or subtype_of_transform == "br": + if do_quick == True: + syn_convergence = "[50x0,1e-6,10]" + else: + syn_convergence = "[50x20,1e-6,10]" + syn_shrink_factors = "2x1" + syn_smoothing_sigmas = "1x0vox" + + syn_stage = [ + "--metric", + syn_metric, + ] + + if multivariate_extras is not None: + for kk in range(len(multivariate_extras)): + syn_stage.append("--metric") + metricname = multivariate_extras[kk][0] + metricfixed = get_pointer_string( + multivariate_extras[kk][1] + ) + metricmov = get_pointer_string( + multivariate_extras[kk][2] + ) + metricWeight = multivariate_extras[kk][3] + metricSampling = multivariate_extras[kk][4] + metricString = "%s[%s,%s,%s,%s]" % ( + metricname, + metricfixed, + metricmov, + metricWeight, + metricSampling, + ) + syn_stage.append(metricString) + + syn_stage.append("--convergence") + syn_stage.append(syn_convergence) + syn_stage.append("--shrink-factors") + syn_stage.append(syn_shrink_factors) + syn_stage.append("--smoothing-sigmas") + syn_stage.append(syn_smoothing_sigmas) + + if ( + subtype_of_transform == "b" + or subtype_of_transform == "br" + or subtype_of_transform == "bo" + ): + syn_stage.insert(0, "BSplineSyN[0.1," + str(spline_distance) + ",0,3]") + syn_stage.insert(0, "--transform") + + if ( + subtype_of_transform == "s" + or subtype_of_transform == "sr" + or subtype_of_transform == "so" + ): + syn_stage.insert(0, "SyN[0.1,3,0]") + syn_stage.insert(0, "--transform") + + args = [ + "-d", + str(fixed.dimension), + "-r" + ] + initx + [ + "-o", + "[%s,%s,%s]" % (outprefix, wmo, wfo), + ] + + if subtype_of_transform == "r" or subtype_of_transform == "t": + args.append(rigid_stage) + if subtype_of_transform == "a": + args.append(rigid_stage) + args.append(affine_stage) + if subtype_of_transform == "b" or subtype_of_transform == "s": + args.append(rigid_stage) + args.append(affine_stage) + args.append(syn_stage) + if subtype_of_transform == "br" or subtype_of_transform == "sr": + args.append(rigid_stage) + args.append(syn_stage) + if subtype_of_transform == "bo" or subtype_of_transform == "so": + args.append(syn_stage) + + args.append("-x") + args.append(maskopt) + + args = list( + itertools.chain.from_iterable( + itertools.repeat(x, 1) if isinstance(x, str) else x + for x in args + ) + ) + + # ------------------------------------------------------------ + + if random_seed is not None: + args.append("--random-seed") + args.append(random_seed) + + if restrict_transformation is not None: + args.append("-g") + args.append(restrict_transformationchar) + + args.append("--float") + args.append(str(int(singleprecision))) + args.append("--write-composite-transform") + args.append(write_composite_transform * 1) + if verbose: + args.append("-v") + args.append("1") + + processed_args = process_arguments(args) + libfn = get_lib_fn("antsRegistration") + if verbose: + print("antsRegistration " + ' '.join(processed_args)) + reg_exit = libfn(processed_args) + if (reg_exit != 0): + raise RuntimeError(f"Registration failed with error code {reg_exit}") + afffns = glob.glob(outprefix + "*" + "[0-9]GenericAffine.mat") + fwarpfns = glob.glob(outprefix + "*" + "[0-9]Warp.nii.gz") + iwarpfns = glob.glob(outprefix + "*" + "[0-9]InverseWarp.nii.gz") + vfieldfns = glob.glob(outprefix + "*" + "[0-9]VelocityField.nii.gz") + # print(afffns, fwarpfns, iwarpfns) + if len(afffns) == 0: + afffns = "" + if len(fwarpfns) == 0: + fwarpfns = "" + if len(iwarpfns) == 0: + iwarpfns = "" + if len(vfieldfns) == 0: + vfieldfns = "" + + alltx = sorted( + set(glob.glob(outprefix + "*" + "[0-9]*")) + - set(glob.glob(outprefix + "*VelocityField*")) + ) + findinv = np.where( + [re.search("[0-9]InverseWarp.nii.gz", ff) for ff in alltx] + )[0] + findfwd = np.where([re.search("[0-9]Warp.nii.gz", ff) for ff in alltx])[ + 0 + ] + if len(findinv) > 0: + fwdtransforms = list( + reversed( + [ff for idx, ff in enumerate(alltx) if idx != findinv[0]] + ) + ) + invtransforms = [ + ff for idx, ff in enumerate(alltx) if idx != findfwd[0] + ] + else: + fwdtransforms = list(reversed(alltx)) + invtransforms = alltx + + if write_composite_transform: + fwdtransforms = outprefix + "Composite.h5" + invtransforms = outprefix + "InverseComposite.h5" + + if not vfieldfns: + return { + "warpedmovout": warpedmovout.clone(inpixeltype), + "warpedfixout": warpedfixout.clone(inpixeltype), + "fwdtransforms": fwdtransforms, + "invtransforms": invtransforms, + } + else: + return { + "warpedmovout": warpedmovout.clone(inpixeltype), + "warpedfixout": warpedfixout.clone(inpixeltype), + "fwdtransforms": fwdtransforms, + "invtransforms": invtransforms, + "velocityfield": vfieldfns, + } + +def motion_correction( + image, + fixed=None, + type_of_transform="BOLDRigid", + mask=None, + fdOffset=50, + outprefix="", + verbose=False, + **kwargs +): + """ + Correct time-series data for motion. + + ANTsR function: `antsrMotionCalculation` + + Arguments + --------- + image: antsImage, usually ND where D=4. + + fixed: Fixed image to register all timepoints to. If not provided, + mean image is used. + + type_of_transform : string + A linear or non-linear registration type. Mutual information metric and rigid transformation by default. + See ants registration for details. + + mask: mask for image (ND-1). If not provided, estimated from data. + 2023-02-05: a performance change - previously, we estimated a mask + when None is provided and would pass this to the registration. this + impairs performance if the mask estimate is bad. in such a case, we + prefer no mask at all. As such, we no longer pass the mask to the + registration when None is provided. + + fdOffset: offset value to use in framewise displacement calculation + + outprefix : string + output will be named with this prefix plus a numeric extension. + + verbose: boolean + + kwargs: keyword args + extra arguments - these extra arguments will control the details of registration that is performed. see ants registration for more. + + Returns + ------- + dict containing follow key/value pairs: + `motion_corrected`: Moving image warped to space of fixed image. + `motion_parameters`: transforms for each image in the time series. + `FD`: Framewise displacement generalized for arbitrary transformations. + + Notes + ----- + Control extra arguments via kwargs. see ants.registration for details. + + Example + ------- + >>> import ants + >>> fi = ants.image_read(ants.get_ants_data('ch2')) + >>> mytx = ants.motion_correction( fi ) + """ + idim = image.dimension + ishape = image.shape + nTimePoints = ishape[idim - 1] + if fixed is None: + wt = 1.0 / nTimePoints + fixed = ants.slice_image(image, axis=idim - 1, idx=0) * 0 + for k in range(nTimePoints): + temp = ants.slice_image(image, axis=idim - 1, idx=k) + fixed = fixed + ants.iMath(temp,"Normalize") * wt + if mask is None: + mask = ants.get_mask(fixed) + useMask=None + else: + useMask=mask + FD = np.zeros(nTimePoints) + motion_parameters = list() + motion_corrected = list() + centerOfMass = mask.get_center_of_mass() + npts = pow(2, idim - 1) + pointOffsets = np.zeros((npts, idim - 1)) + myrad = np.ones(idim - 1).astype(int).tolist() + mask1vals = np.zeros(int(mask.sum())) + mask1vals[round(len(mask1vals) / 2)] = 1 + mask1 = ants.make_image(mask, mask1vals) + myoffsets = ants.get_neighborhood_in_mask( + mask1, mask1, radius=myrad, spatial_info=True + )["offsets"] + + mycols = list("xy") + if idim - 1 == 3: + mycols = list("xyz") + useinds = list() + for k in range(myoffsets.shape[0]): + if abs(myoffsets[k, :]).sum() == (idim - 2): + useinds.append(k) + myoffsets[k, :] = myoffsets[k, :] * fdOffset / 2.0 + centerOfMass + fdpts = pd.DataFrame(data=myoffsets[useinds, :], columns=mycols) + if verbose: + print("Progress:") + counter = 0 + for k in range(nTimePoints): + mycount = round(k / nTimePoints * 100) + if verbose and mycount == counter: + counter = counter + 10 + print(mycount, end="%.", flush=True) + temp = ants.slice_image(image, axis=idim - 1, idx=k) + temp = ants.iMath(temp, "Normalize") + if temp.numpy().var() > 0: + if outprefix != "": + outprefixloc = outprefix + "_" + str.zfill( str(k), 5 ) + "_" + myreg = registration( + fixed, temp, type_of_transform=type_of_transform, mask=useMask, + outprefix=outprefixloc, **kwargs + ) + else: + myreg = registration( + fixed, temp, type_of_transform=type_of_transform, mask=useMask, **kwargs + ) + fdptsTxI = ants.apply_transforms_to_points( + idim - 1, fdpts, myreg["fwdtransforms"] + ) + if k > 0 and motion_parameters[k - 1] != "NA": + fdptsTxIminus1 = ants.apply_transforms_to_points( + idim - 1, fdpts, motion_parameters[k - 1] + ) + else: + fdptsTxIminus1 = fdptsTxI + # take the absolute value, then the mean across columns, then the sum + FD[k] = (fdptsTxIminus1 - fdptsTxI).abs().mean().sum() + motion_parameters.append(myreg["fwdtransforms"]) + mywarped = ants.apply_transforms( fixed, + ants.slice_image(image, axis=idim - 1, idx=k), + myreg["fwdtransforms"] ) + motion_corrected.append(mywarped) + else: + motion_parameters.append("NA") + motion_corrected.append(temp) + + if verbose: + print("Done") + return { + "motion_corrected": ants.list_to_ndimage(image, motion_corrected), + "motion_parameters": motion_parameters, + "FD": FD, + } + +def label_image_registration(fixed_label_images, + moving_label_images, + fixed_intensity_images=None, + moving_intensity_images=None, + fixed_mask=None, + moving_mask=None, + type_of_linear_transform='affine', + type_of_deformable_transform='antsRegistrationSyNQuick[so]', + label_image_weighting=1.0, + output_prefix='', + random_seed=None, + verbose=False): + + """ + Perform pairwise registration using fixed and moving sets of label + images (and, optionally, sets of corresponding intensity images). + + Arguments + --------- + fixed_label_images : single or list of ANTsImage + A single (or set of) fixed label image(s). + + moving_label_images : single or list of ANTsImage + A single (or set of) moving label image(s). + + fixed_intensity_images : single or list of ANTsImage + Optional---a single (or set of) fixed intensity image(s). + + moving_intensity_images : single or list of ANTsImage + Optional---a single (or set of) moving intensity image(s). + + fixed_mask : ANTsImage + Defines region for similarity metric calculation in the space + of the fixed image. + + moving_mask : ANTsImage + Defines region for similarity metric calculation in the space + of the moving image. + + type_of_linear_transform : string + Use label images with the centers of mass to a calculate linear + transform of type 'rigid', 'similarity', or 'affine'. + + type_of_deformable_transform : string + Only works with deformable-only transforms, specifically the family + of antsRegistrationSyN*[so] or antsRegistrationSyN*[bo] transforms. + See 'type_of_transform' in ants.registration. Additionally, one can + use a list to pass a more tailored deformably-only transform + optimization using SyN or BSplineSyN transforms. The order of + parameters in the list would be 1) transform specification, i.e. + "SyN" or "BSplineSyN", 2) gradient (real), 3) intensity metric (string), + 4) intensity metric parameter (real), 5) convergence iterations per level + (tuple) 6) smoothing factors per level (tuple), 7) shrink factors per level + (tuple). An example would type_of_deformable_transform = ["SyN", 0.2, "CC", + 4, (100,50,10), (2,1,0), (4,2,1)]. + + label_image_weighting : float or list of floats + Relative weighting for the label images. + + output_prefix : string + Define the output prefix for the filenames of the output transform + files. + + random_seed : integer + Definition for deformable registration. + + verbose : boolean + Print progress to the screen. + + Returns + ------- + Set of transforms definining the mapping to/from the fixed image domain + to the moving image domain. + + Example + ------- + >>> import ants + >>> + >>> r16 = ants.image_read(ants.get_ants_data('r16')) + >>> r16_seg1 = ants.threshold_image(r16, "Kmeans", 3) - 1 + >>> r16_seg2 = ants.threshold_image(r16, "Kmeans", 5) - 1 + >>> r64 = ants.image_read(ants.get_ants_data('r64')) + >>> r64_seg1 = ants.threshold_image(r64, "Kmeans", 3) - 1 + >>> r64_seg2 = ants.threshold_image(r64, "Kmeans", 5) - 1 + >>> reg = ants.label_image_registration([r16_seg1, r16_seg2], + [r64_seg1, r64_seg2], + fixed_intensity_images=r16, + moving_intensity_images=r64, + type_of_linear_transform='affine', + type_of_deformable_transform='antsRegistrationSyNQuick[bo]', + label_image_weighting=[1.0, 2.0], + verbose=True) + """ + + # Perform validation check on the input + + if isinstance(fixed_label_images, ants.ANTsImage): + fixed_label_images = [ants.image_clone(fixed_label_images)] + if isinstance(moving_label_images, ants.ANTsImage): + moving_label_images = [ants.image_clone(moving_label_images)] + + if len(fixed_label_images) != len(moving_label_images): + raise ValueError("The number of fixed and moving label images do not match.") + + if fixed_intensity_images is not None or moving_intensity_images is not None: + if isinstance(fixed_intensity_images, ants.ANTsImage): + fixed_intensity_images = [ants.image_clone(fixed_intensity_images)] + if isinstance(moving_intensity_images, ants.ANTsImage): + moving_intensity_images = [ants.image_clone(moving_intensity_images)] + if len(fixed_intensity_images) != len(moving_intensity_images): + raise ValueError("The number of fixed and moving intensity images do not match.") + + label_image_weights = list() + if isinstance(label_image_weighting, (int, float)): + label_image_weights = [label_image_weighting] * len(fixed_label_images) + else: + label_image_weights = tuple(label_image_weighting) + if len(fixed_label_images) != len(label_image_weights): + raise ValueError("The length of label_image_weights must" + + "match the number of label image pairs.") + + image_dimension = fixed_label_images[0].dimension + + if output_prefix == "" or output_prefix is None or len(output_prefix) == 0: + output_prefix = mktemp() + + allowable_linear_transforms = ['rigid', 'similarity', 'affine'] + if not type_of_linear_transform in allowable_linear_transforms: + raise ValueError("Unrecognized linear transform.") + + do_deformable = True + if type_of_deformable_transform is None or len(type_of_deformable_transform) == 0: + do_deformable = False + + common_label_ids = list() + total_number_of_labels = 0 + for i in range(len(fixed_label_images)): + fixed_label_geoms = ants.label_geometry_measures(fixed_label_images[i]) + fixed_label_ids = np.array(fixed_label_geoms['Label']) + moving_label_geoms = ants.label_geometry_measures(moving_label_images[i]) + moving_label_ids = np.array(moving_label_geoms['Label']) + common_label_ids.append(np.intersect1d(moving_label_ids, fixed_label_ids)) + total_number_of_labels += len(common_label_ids[i]) + if verbose: + print("Common label ids for image pair ", str(i), ": ", common_label_ids[i]) + if len(common_label_ids[i]) == 0: + raise ValueError("No common labels for image pair " + str(i)) + + if verbose: + print("Total number of labels: " + str(total_number_of_labels)) + + ############################## + # + # Linear transform + # + ############################## + + linear_xfrm = None + if type_of_linear_transform is not None: + + if verbose: + print("\n\nComputing linear transform.\n") + + if total_number_of_labels < 3: + raise ValueError(" Number of labels must be >= 3.") + + fixed_centers_of_mass = np.zeros((total_number_of_labels, image_dimension)) + moving_centers_of_mass = np.zeros((total_number_of_labels, image_dimension)) + deformable_multivariate_extras = list() + + count = 0 + for i in range(len(common_label_ids)): + for j in range(len(common_label_ids[i])): + label = common_label_ids[i][j] + if verbose: + print(" Finding centers of mass for image pair " + str(i) + ", label " + str(label)) + fixed_single_label_image = ants.threshold_image(fixed_label_images[i], label, label, 1, 0) + fixed_centers_of_mass[count, :] = ants.get_center_of_mass(fixed_single_label_image) + moving_single_label_image = ants.threshold_image(moving_label_images[i], label, label, 1, 0) + moving_centers_of_mass[count, :] = ants.get_center_of_mass(moving_single_label_image) + count += 1 + if do_deformable: + deformable_multivariate_extras.append(["MSQ", fixed_single_label_image, + moving_single_label_image, + label_image_weights[i], 0]) + + linear_xfrm = ants.fit_transform_to_paired_points(moving_centers_of_mass, + fixed_centers_of_mass, + transform_type=type_of_linear_transform, + verbose=verbose) + + linear_xfrm_file = output_prefix + "0GenericAffine.mat" + ants.write_transform(linear_xfrm, linear_xfrm_file) + + ############################## + # + # Deformable transform + # + ############################## + + if do_deformable: + + if verbose: + print("\n\nComputing deformable transform using images.\n") + + intensity_metric = "CC" + intensity_metric_parameter = 2 + syn_shrink_factors = "8x4x2x1" + syn_smoothing_sigmas = "3x2x1x0vox" + syn_convergence = "[100x70x50x20,1e-6,10]" + spline_distance = 26 + gradient_step = 0.1 + syn_transform = "SyN" + + syn_stage = list() + + if isinstance(type_of_deformable_transform, list): + + if (len(type_of_deformable_transform) != 7 or + not isinstance(type_of_deformable_transform[0], str) or + not isinstance(type_of_deformable_transform[1], float) or + not isinstance(type_of_deformable_transform[2], str) or + not isinstance(type_of_deformable_transform[3], int) or + not isinstance(type_of_deformable_transform[4], tuple) or + not isinstance(type_of_deformable_transform[5], tuple) or + not isinstance(type_of_deformable_transform[6], tuple)): + raise ValueError("Incorrect specification for type_of_deformable_transform. See help menu.") + + syn_transform = type_of_deformable_transform[0] + gradient_step = type_of_deformable_transform[1] + intensity_metric = type_of_deformable_transform[2] + intensity_metric_parameter = type_of_deformable_transform[3] + + t = type_of_deformable_transform[4] + tstr = ''.join(map(lambda x: str(x) + 'x', t[:len(t)-1])) + str(t[len(t)-1]) + syn_convergence = "[" + tstr + ",1e-6,10]" + + t = type_of_deformable_transform[5] + tstr = ''.join(map(lambda x: str(x) + 'x', t[:len(t)-1])) + str(t[len(t)-1]) + syn_smoothing_sigmas = tstr + "vox" + + t = type_of_deformable_transform[6] + syn_shrink_factors = ''.join(map(lambda x: str(x) + 'x', t[:len(t)-1])) + str(t[len(t)-1]) + + else: + + do_quick = False + if "Quick" in type_of_deformable_transform: + do_quick = True + elif "Repro" in type_of_deformable_transform: + random_seed = str(1) + + if "[" in type_of_deformable_transform and "]" in type_of_deformable_transform: + subtype_of_deformable_transform = type_of_deformable_transform.split("[")[1].split("]")[0] + if not ('bo' in subtype_of_deformable_transform or 'so' in subtype_of_deformable_transform): + raise ValueError("Only 'so' or 'bo' transforms are available.") + else: + if 'bo' in subtype_of_deformable_transform: + syn_transform = "BSplineSyN" + if "," in subtype_of_deformable_transform: + subtype_of_deformable_transform_args = subtype_of_deformable_transform.split(",") + subtype_of_deformable_transform = subtype_of_deformable_transform_args[0] + intensity_metric_parameter = subtype_of_deformable_transform_args[1] + if len(subtype_of_deformable_transform_args) > 2: + spline_distance = subtype_of_deformable_transform_args[2] + + if do_quick: + intensity_metric = "MI" + if intensity_metric_parameter is None: + intensity_metric_parameter = 32 + syn_convergence = "[100x70x50x0,1e-6,10]" + + if fixed_intensity_images is not None and len(fixed_intensity_images) > 0: + for i in range(len(fixed_intensity_images)): + syn_stage.append("--metric") + metric_string = "%s[%s,%s,%s,%s]" % ( + intensity_metric, + get_pointer_string(fixed_intensity_images[i]), + get_pointer_string(moving_intensity_images[i]), + 1.0, intensity_metric_parameter) + syn_stage.append(metric_string) + + for kk in range(len(deformable_multivariate_extras)): + syn_stage.append("--metric") + metricString = "%s[%s,%s,%s,%s]" % ( + "MSQ", + get_pointer_string(deformable_multivariate_extras[kk][1]), + get_pointer_string(deformable_multivariate_extras[kk][2]), + deformable_multivariate_extras[kk][3], 0.0) + syn_stage.append(metricString) + + syn_stage.append("--convergence") + syn_stage.append(syn_convergence) + syn_stage.append("--shrink-factors") + syn_stage.append(syn_shrink_factors) + syn_stage.append("--smoothing-sigmas") + syn_stage.append(syn_smoothing_sigmas) + + if syn_transform == "SyN": + syn_stage.insert(0, "SyN[" + str(gradient_step) + ",3,0]") + else: + syn_stage.insert(0, "BSplineSyN[" + str(gradient_step) + "," + str(spline_distance) + ",0,3]") + syn_stage.insert(0, "--transform") + + args = None + if linear_xfrm is None: + args = ["-d", str(image_dimension), + "-o", output_prefix] + else: + args = ["-d", str(image_dimension), + "-r", linear_xfrm_file, + "-o", output_prefix] + args.append(syn_stage) + + fixed_mask_string = 'NA' + if fixed_mask is not None: + fixed_mask_binary = fixed_mask != 0 + fixed_mask_string = get_pointer_string(fixed_mask_binary) + + moving_mask_string = 'NA' + if moving_mask is not None: + moving_mask_binary = moving_mask != 0 + moving_mask_string = get_pointer_string(moving_mask_binary) + + mask_option = "[%s,%s]" % (fixed_mask_string, moving_mask_string) + + args.append("-x") + args.append(mask_option) + + args = list(itertools.chain.from_iterable( + itertools.repeat(x, 1) + if isinstance(x, str) + else x for x in args)) + + args.append("--float") + args.append("1") + + if random_seed is not None: + args.append("--random-seed") + args.append(random_seed) + + if verbose: + args.append("-v") + args.append("1") + + processed_args = process_arguments(args) + if verbose: + print("antsRegistration " + ' '.join(processed_args)) + + libfn = get_lib_fn("antsRegistration") + deformable_registration_exit_error = libfn(processed_args) + + if deformable_registration_exit_error != 0: + raise RuntimeError(f"Registration failed with error code {deformable_registration_exit_error}") + + all_xfrms = sorted(set(glob.glob(output_prefix + "*" + "[0-9]*"))) + + find_inverse_warps = np.where([re.search("[0-9]InverseWarp.nii.gz", ff) for ff in all_xfrms])[0] + find_forward_warps = np.where([re.search("[0-9]Warp.nii.gz", ff) for ff in all_xfrms])[0] + + if len(find_inverse_warps) > 0: + fwdtransforms = [all_xfrms[find_forward_warps[0]], linear_xfrm_file] + invtransforms = [linear_xfrm_file, all_xfrms[find_inverse_warps[0]]] + else: + fwdtransforms = [linear_xfrm_file] + invtransforms = [linear_xfrm_file] + + if verbose: + print("\n\nResulting transforms") + print(" fwdtransforms: ", fwdtransforms) + print(" invtransforms: ", invtransforms) + + return { + "fwdtransforms": fwdtransforms, + "invtransforms": invtransforms, + } + + diff --git a/MindEyeV2/antspy/ants/registration/simulate_displacement_field.py b/MindEyeV2/antspy/ants/registration/simulate_displacement_field.py new file mode 100644 index 0000000000000000000000000000000000000000..6fc8ab2b38dd1881337c93e1098a24a7e2c1f958 --- /dev/null +++ b/MindEyeV2/antspy/ants/registration/simulate_displacement_field.py @@ -0,0 +1,90 @@ +__all__ = ["simulate_displacement_field"] + +import numpy as np + + +import ants +from ants.internal import get_lib_fn + + + +def simulate_displacement_field(domain_image, + field_type="bspline", + number_of_random_points=1000, + sd_noise=10.0, + enforce_stationary_boundary=True, + number_of_fitting_levels=4, + mesh_size=1, + sd_smoothing=4.0): + """ + simulate displacement field using either b-spline or exponential transform + + ANTsR function: `simulateDisplacementField` + + Arguments + --------- + domain_image : ANTsImage + Domain image + + field_type : string + Either "bspline" or "exponential". + + number_of_random_points : integer + Number of displacement points. + + sd_noise : float + Standard deviation of the displacement field noise. + + enforce_stationary_boundary : boolean + Determines fixed boundary conditions. + + number_of_fitting_levels : integer + Number of fitting levels (b-spline only). + + mesh_size : integer or n-D tuple + Determines fitting resolution at base level (b-spline only). + + sd_smoothing : float + Standard deviation of the Gaussian smoothing in mm (exponential only). + + Returns + ------- + ANTs vector image. + + Example + ------- + >>> import ants + >>> domain = ants.image_read( ants.get_ants_data('r16')) + >>> exp_field = ants.simulate_displacement_field(domain, field_type="exponential") + >>> bsp_field = ants.simulate_displacement_field(domain, field_type="bspline") + >>> bsp_xfrm = ants.transform_from_displacement_field(bsp_field * 3) + >>> domain_warped = ants.apply_ants_transform_to_image(bsp_xfrm, domain, domain) + """ + + image_dimension = domain_image.dimension + + if field_type == 'bspline': + if isinstance(mesh_size, int) == False and len(mesh_size) != image_dimension: + raise ValueError("Incorrect specification for mesh_size.") + + spline_order = 3 + number_of_control_points = mesh_size + spline_order + + if isinstance(number_of_control_points, int) == True: + number_of_control_points = np.repeat(number_of_control_points, image_dimension) + + libfn = get_lib_fn("simulateBsplineDisplacementField%iD" % image_dimension) + field = libfn(domain_image.pointer, number_of_random_points, sd_noise, + enforce_stationary_boundary, number_of_fitting_levels, number_of_control_points) + bspline_field = ants.from_pointer(field).clone('float') + return bspline_field + + elif field_type == 'exponential': + libfn = get_lib_fn("simulateExponentialDisplacementField%iD" % image_dimension) + field = libfn(domain_image.pointer, number_of_random_points, sd_noise, + enforce_stationary_boundary, sd_smoothing) + exp_field = ants.from_pointer(field).clone('float') + return exp_field + + else: + raise ValueError("Unrecognized field type.") diff --git a/MindEyeV2/src/slurms/458689.out b/MindEyeV2/src/slurms/458689.out new file mode 100644 index 0000000000000000000000000000000000000000..82683a3b77ca4d3a9eaa40c7504f0d997cd102cc --- /dev/null +++ b/MindEyeV2/src/slurms/458689.out @@ -0,0 +1,4 @@ +MASTER_ADDR=ip-10-0-158-103 +MASTER_PORT=11437 +WORLD_SIZE=1 +model_name=semantic_cluster_1.2_average_after_wd-2_no_prior_multi diff --git a/MindEyeV2/src/slurms/458690.err b/MindEyeV2/src/slurms/458690.err new file mode 100644 index 0000000000000000000000000000000000000000..7b1aec4ca829bc1c8b35bb082f6cf2c74bd48f8f --- /dev/null +++ b/MindEyeV2/src/slurms/458690.err @@ -0,0 +1,77 @@ +[NbConvertApp] Converting notebook Untitled1.ipynb to python +[NbConvertApp] Writing 43585 bytes to Untitled1.py +wandb: Currently logged in as: ckadirt. Use `wandb login --relogin` to force relogin +wandb: wandb version 0.17.4 is available! To upgrade, please run: +wandb: $ pip install wandb --upgrade +wandb: Tracking run with wandb version 0.17.1 +wandb: Run data is saved locally in /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi +wandb: Run `wandb offline` to turn off syncing. +wandb: Syncing run semantic_cluster_1.2_average_after_wd-2_no_prior_multi +wandb: ⭐️ View project at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster_0.2 +wandb: 🚀 View run at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster_0.2/runs/semantic_cluster_1.2_average_after_wd-2_no_prior_multi + 0%| | 0/150 [00:00 + if plotting: + ^^^^^^^^ +NameError: name 'plotting' is not defined +Traceback (most recent call last): + File "/weka/proj-fmri/ckadirt/MindEyeV2/src/enhanced_recon_inference_old.py", line 100, in + all_recons = torch.load(f"evals/{model_name}/{model_name}_all_recons.pt") # these are the unrefined MindEye2 recons! + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/torch/serialization.py", line 986, in load + with _open_file_like(f, 'rb') as opened_file: + ^^^^^^^^^^^^^^^^^^^^^^^^ + File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/torch/serialization.py", line 435, in _open_file_like + return _open_file(name_or_buffer, mode) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/torch/serialization.py", line 416, in __init__ + super().__init__(open(name, mode)) + ^^^^^^^^^^^^^^^^ +FileNotFoundError: [Errno 2] No such file or directory: 'evals/final_subj01_pretrained_3sess_24bs/final_subj01_pretrained_3sess_24bs_all_recons.pt' diff --git a/MindEyeV2/src/slurms/534074.out b/MindEyeV2/src/slurms/534074.out new file mode 100644 index 0000000000000000000000000000000000000000..ac5d610022bb68426d26c9ecdce6a5cb4e9e137f --- /dev/null +++ b/MindEyeV2/src/slurms/534074.out @@ -0,0 +1,44 @@ +MASTER_ADDR=ip-10-0-150-188 +MASTER_PORT=13997 +final_subj01_pretrained_3sess_24bs +new_sessions +device: cuda +num_voxels for subj01: 15724 +/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar +Loaded test dl for subj1! + +0 3000 3000 1000 +param counts: +83,653,863 total +0 trainable +param counts: +64,409,600 total +64,409,600 trainable +param counts: +1,903,020,028 total +1,903,020,028 trainable +param counts: +1,967,429,628 total +1,967,429,628 trainable +param counts: +259,865,216 total +259,865,200 trainable +param counts: +2,227,294,844 total +2,227,294,828 trainable + +---loading /weka/proj-fmri/ckadirt/MindEyeV2/train_logs/final_subj01_pretrained_3sess_24bs/last.pth ckpt--- + +[2024-11-07 02:45:52,200] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect) +Processing zero checkpoint '/weka/proj-fmri/ckadirt/MindEyeV2/train_logs/final_subj01_pretrained_3sess_24bs/last' +Detected checkpoint of type zero stage ZeroStageEnum.gradients, world_size: 8 +Parsing checkpoint created by deepspeed==0.12.2 +Reconstructed Frozen fp32 state dict with 1 params 16 elements +Reconstructed fp32 state dict with 230 params 2227294828 elements +ckpt loaded! +Initialized embedder #0: FrozenOpenCLIPImageEmbedder with 1909889025 params. Trainable: False +Initialized embedder #1: ConcatTimestepEmbedderND with 0 params. Trainable: False +Initialized embedder #2: ConcatTimestepEmbedderND with 0 params. Trainable: False +vector_suffix torch.Size([1, 1024]) +['a group of people sitting around a table.'] +device: cuda diff --git a/MindEyeV2/src/slurms/534079.err b/MindEyeV2/src/slurms/534079.err new file mode 100644 index 0000000000000000000000000000000000000000..dd6e2daa8072955d4df5cc45a3ceeda8f7ec9a0a --- /dev/null +++ b/MindEyeV2/src/slurms/534079.err @@ -0,0 +1,4061 @@ +[NbConvertApp] Converting notebook enhanced_recon_inference_old.ipynb to python +[NbConvertApp] Writing 13091 bytes to enhanced_recon_inference_old.py +[NbConvertApp] Converting notebook recon_inference_old.ipynb to python +[NbConvertApp] Writing 17081 bytes to recon_inference_old.py +/admin/home-ckadirt/fmri/lib/python3.11/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`. + warnings.warn( +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now. +WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now. + 0%| | 0/1000 [00:00 + "train_url": train_url, + ^^^^^^^^^ +NameError: name 'train_url' is not defined. Did you mean: 'train_dl'? diff --git a/MindEyeV2/src/slurms/544387.out b/MindEyeV2/src/slurms/544387.out new file mode 100644 index 0000000000000000000000000000000000000000..42b1826d939b4297f5529044c3b9c520d14e9dbf --- /dev/null +++ b/MindEyeV2/src/slurms/544387.out @@ -0,0 +1,54 @@ +MASTER_ADDR=ip-10-0-130-125 +MASTER_PORT=15229 +WORLD_SIZE=1 +model_name=augmented_image_one +LOCAL RANK 0 +PID of this process = 1825637 +device: cuda +Distributed environment: DistributedType.NO +Num processes: 1 +Process index: 0 +Local process index: 0 +Device: cuda + +Mixed precision type: fp16 + +distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16 +subj_list [1] num_sessions 15 +dividing batch size by subj_list, which will then be concatenated across subj during training... +Training with 15 sessions +Loaded all subj train dls and betas! + +Loaded all subj train dls and betas! + +Loaded test dl for subj1! + +batch_size = 21 num_iterations_per_epoch = 205 num_samples_per_epoch = 4323 +param counts: +712,785,920 total +712,785,920 trainable +param counts: +712,785,920 total +712,785,920 trainable +torch.Size([2, 1, 174019]) torch.Size([2, 1, 4096]) +param counts: +1,887,861,400 total +1,887,861,400 trainable +param counts: +2,600,647,320 total +2,600,647,320 trainable +b.shape torch.Size([2, 1, 4096]) +torch.Size([2, 256, 1664]) torch.Size([2, 256, 1664]) torch.Size([1]) torch.Size([1]) +param counts: +259,865,216 total +259,865,200 trainable +param counts: +2,860,512,536 total +2,860,512,520 trainable +total_steps 16400 + +Done with model preparations! +param counts: +2,860,512,536 total +2,860,512,520 trainable +wandb mindeye run augmented_image_one diff --git a/MindEyeV2/src/slurms/544389.err b/MindEyeV2/src/slurms/544389.err new file mode 100644 index 0000000000000000000000000000000000000000..7c414a06d93c816fd5e93bd3b3caddfb5b729eef --- /dev/null +++ b/MindEyeV2/src/slurms/544389.err @@ -0,0 +1,26 @@ +[NbConvertApp] Converting notebook TrainB5k.ipynb to python +[NbConvertApp] Writing 52131 bytes to TrainB5k.py +wandb: Currently logged in as: ckadirt. Use `wandb login --relogin` to force relogin +wandb: wandb version 0.19.0 is available! To upgrade, please run: +wandb: $ pip install wandb --upgrade +wandb: Tracking run with wandb version 0.17.1 +wandb: Run data is saved locally in /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20241206_230633-augmented_image_one +wandb: Run `wandb offline` to turn off syncing. +wandb: Syncing run augmented_image_one +wandb: ⭐️ View project at https://stability.wandb.io/ckadirt/mindeye +wandb: 🚀 View run at https://stability.wandb.io/ckadirt/mindeye/runs/augmented_image_one + 0%| | 0/80 [00:00 + accelerator.backward(loss) + File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/accelerate/accelerator.py", line 1987, in backward + self.scaler.scale(loss).backward(**kwargs) + File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward + torch.autograd.backward( + File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward + Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass +torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.50 GiB. GPU 0 has a total capacty of 79.11 GiB of which 1006.94 MiB is free. Including non-PyTorch memory, this process has 78.12 GiB memory in use. Of the allocated memory 61.43 GiB is allocated by PyTorch, and 15.89 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF +wandb: - 0.071 MB of 0.071 MB uploaded wandb: \ 0.071 MB of 0.071 MB uploaded wandb: | 0.071 MB of 0.071 MB uploaded wandb: / 0.077 MB of 0.085 MB uploaded wandb: - 0.085 MB of 0.085 MB uploaded wandb: 🚀 View run augmented_image_one at: https://stability.wandb.io/ckadirt/mindeye/runs/augmented_image_one +wandb: ⭐️ View project at: https://stability.wandb.io/ckadirt/mindeye +wandb: Synced 5 W&B file(s), 0 media file(s), 3 artifact file(s) and 1 other file(s) +wandb: Find logs at: ./wandb/run-20241206_230633-augmented_image_one/logs diff --git a/MindEyeV2/src/slurms/544389.out b/MindEyeV2/src/slurms/544389.out new file mode 100644 index 0000000000000000000000000000000000000000..df062d8ff0a2beca5c675a56ccc4f784da51f5ec --- /dev/null +++ b/MindEyeV2/src/slurms/544389.out @@ -0,0 +1,59 @@ +MASTER_ADDR=ip-10-0-130-125 +MASTER_PORT=11619 +WORLD_SIZE=1 +model_name=augmented_image_one +LOCAL RANK 0 +PID of this process = 1826787 +device: cuda +Distributed environment: DistributedType.NO +Num processes: 1 +Process index: 0 +Local process index: 0 +Device: cuda + +Mixed precision type: fp16 + +distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16 +subj_list [1] num_sessions 15 +dividing batch size by subj_list, which will then be concatenated across subj during training... +Training with 15 sessions +Loaded all subj train dls and betas! + +Loaded all subj train dls and betas! + +Loaded test dl for subj1! + +batch_size = 21 num_iterations_per_epoch = 205 num_samples_per_epoch = 4323 +param counts: +712,785,920 total +712,785,920 trainable +param counts: +712,785,920 total +712,785,920 trainable +torch.Size([2, 1, 174019]) torch.Size([2, 1, 4096]) +param counts: +1,887,861,400 total +1,887,861,400 trainable +param counts: +2,600,647,320 total +2,600,647,320 trainable +b.shape torch.Size([2, 1, 4096]) +torch.Size([2, 256, 1664]) torch.Size([2, 256, 1664]) torch.Size([1]) torch.Size([1]) +param counts: +259,865,216 total +259,865,200 trainable +param counts: +2,860,512,536 total +2,860,512,520 trainable +total_steps 16400 + +Done with model preparations! +param counts: +2,860,512,536 total +2,860,512,520 trainable +wandb mindeye run augmented_image_one +wandb_config: + {'model_name': 'augmented_image_one', 'global_batch_size': '21', 'batch_size': 21, 'num_epochs': 80, 'num_sessions': 15, 'num_params': 2860512520, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': True, 'max_lr': 0.0003, 'mixup_pct': 0.33, 'num_samples_per_epoch': 4323, 'num_test': 480, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1} +wandb_id: augmented_image_one +torch.Size([21, 174019]) torch.Size([21, 3, 224, 224]) torch.Size([21]) torch.Size([21]) +augmented_image_one starting with epoch 0 / 80 diff --git a/MindEyeV2/src/slurms/544390.err b/MindEyeV2/src/slurms/544390.err new file mode 100644 index 0000000000000000000000000000000000000000..63ede4cc0ee2f0b737af7c9b865e1a0a567c9984 --- /dev/null +++ b/MindEyeV2/src/slurms/544390.err @@ -0,0 +1,67 @@ +[NbConvertApp] Converting notebook TrainB5k.ipynb to python +[NbConvertApp] Writing 52131 bytes to TrainB5k.py +wandb: Currently logged in as: ckadirt. Use `wandb login --relogin` to force relogin +wandb: wandb version 0.19.0 is available! To upgrade, please run: +wandb: $ pip install wandb --upgrade +wandb: Tracking run with wandb version 0.17.1 +wandb: Run data is saved locally in /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20241206_230853-augmented_image_one_ +wandb: Run `wandb offline` to turn off syncing. +wandb: Syncing run augmented_image_one_ +wandb: ⭐️ View project at https://stability.wandb.io/ckadirt/mindeye +wandb: 🚀 View run at https://stability.wandb.io/ckadirt/mindeye/runs/augmented_image_one_ + 0%| | 0/80 [00:00 + aaa +NameError: name 'aaa' is not defined +wandb: - 0.071 MB of 0.071 MB uploaded wandb: \ 0.071 MB of 0.071 MB uploaded wandb: | 0.071 MB of 0.071 MB uploaded wandb: / 0.071 MB of 0.071 MB uploaded wandb: - 0.079 MB of 0.091 MB uploaded (0.004 MB deduped) wandb: \ 0.091 MB of 0.091 MB uploaded (0.004 MB deduped) wandb: +wandb: Run history: +wandb: test/blurry_pixcorr ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: test/loss █▄▄▃▃▃▂▂▃▂▃▂▂▂▃▅▂▃▃▁▂▂▂▂▃▄▂▃▃▂▂▄▃▃▂▃▃▂▃▃ +wandb: test/loss_clip_total █▁▄▄▃▄▅▃▄▅▅▄▆▅▅▆▅▆▆▆▆▆▆▆▆▆▆▇▇▆▇▇▇▇▇▇▇▇██ +wandb: test/loss_prior █▅▄▃▃▃▂▂▃▂▃▂▂▂▃▅▂▃▃▁▂▂▂▁▃▄▂▃▃▂▂▄▃▃▂▂▃▂▃▃ +wandb: test/num_steps ▁▁▁▂▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▆▇▇▇▇▇███ +wandb: test/recon_cossim ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: test/recon_mse ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: test/test_bwd_pct_correct ▁▅▆▅▃▂▄▄▆▆▄▇▄▅▆▆▆▆▇▇▇▇▇▇█▇▆▆▇█▆▇█▇█▇▇▇▇▇ +wandb: test/test_fwd_pct_correct ▃▂▅▁▂▂▃▃▆▅▄▂▁▅▅▇▅▃▅▄▄▅▄▄▅▃▅█▅▇▅▆▆▆▄▆▅▅▅▅ +wandb: train/blurry_pixcorr ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: train/bwd_pct_correct ▁▃▆▆▆▆▆▆▆▆▆▆▆███████████████████████████ +wandb: train/fwd_pct_correct ▁▄▆▆▆▆▆▆▆▆▆▆▆███████████████████████████ +wandb: train/loss █▄▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: train/loss_blurry_cont_total ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: train/loss_blurry_total ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: train/loss_clip_total █▅▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: train/loss_prior █▄▃▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: train/lr ▅███████▇▇▇▇▇▆▆▆▆▅▅▅▅▄▄▄▃▃▃▃▂▂▂▂▂▁▁▁▁▁▁▁ +wandb: train/num_steps ▁▁▁▂▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▆▇▇▇▇▇███ +wandb: train/recon_cossim ▁▅▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇████████████████ +wandb: train/recon_mse █▄▃▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ +wandb: +wandb: Run summary: +wandb: test/blurry_pixcorr 0.0 +wandb: test/loss 19.72851 +wandb: test/loss_clip_total 6.11602 +wandb: test/loss_prior 0.45375 +wandb: test/num_steps 80 +wandb: test/recon_cossim 0.0 +wandb: test/recon_mse 0.0 +wandb: test/test_bwd_pct_correct 0.02917 +wandb: test/test_fwd_pct_correct 0.025 +wandb: train/blurry_pixcorr 0.0 +wandb: train/bwd_pct_correct 1.0 +wandb: train/fwd_pct_correct 1.0 +wandb: train/loss 8.35443 +wandb: train/loss_blurry_cont_total 0.0 +wandb: train/loss_blurry_total 0.0 +wandb: train/loss_clip_total 1e-05 +wandb: train/loss_prior 0.27848 +wandb: train/lr 0.0 +wandb: train/num_steps 16400 +wandb: train/recon_cossim 0.79336 +wandb: train/recon_mse 0.27848 +wandb: +wandb: 🚀 View run augmented_image_one_ at: https://stability.wandb.io/ckadirt/mindeye/runs/augmented_image_one_ +wandb: ⭐️ View project at: https://stability.wandb.io/ckadirt/mindeye +wandb: Synced 5 W&B file(s), 0 media file(s), 3 artifact file(s) and 1 other file(s) +wandb: Find logs at: ./wandb/run-20241206_230853-augmented_image_one_/logs diff --git a/MindEyeV2/src/slurms/544492.out b/MindEyeV2/src/slurms/544492.out new file mode 100644 index 0000000000000000000000000000000000000000..919736f4febfd16f1e42c1ea2d3ffffcdd0791c7 --- /dev/null +++ b/MindEyeV2/src/slurms/544492.out @@ -0,0 +1,62 @@ +MASTER_ADDR=ip-10-0-172-177 +MASTER_PORT=12285 +WORLD_SIZE=1 +model_name=bold5k_v1 +LOCAL RANK 0 +PID of this process = 1733288 +device: cuda +Distributed environment: DistributedType.NO +Num processes: 1 +Process index: 0 +Local process index: 0 +Device: cuda + +Mixed precision type: fp16 + +distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16 +subj_list [1] num_sessions 15 +dividing batch size by subj_list, which will then be concatenated across subj during training... +Training with 15 sessions +Loaded all subj train dls and betas! + +Loaded all subj train dls and betas! + +Loaded test dl for subj1! + +batch_size = 21 num_iterations_per_epoch = 205 num_samples_per_epoch = 4323 +param counts: +6,905,856 total +6,905,856 trainable +param counts: +6,905,856 total +6,905,856 trainable +torch.Size([2, 1, 1685]) torch.Size([2, 1, 4096]) +param counts: +1,887,861,400 total +1,887,861,400 trainable +param counts: +1,894,767,256 total +1,894,767,256 trainable +b.shape torch.Size([2, 1, 4096]) +torch.Size([2, 256, 1664]) torch.Size([2, 256, 1664]) torch.Size([1]) torch.Size([1]) +param counts: +259,865,216 total +259,865,200 trainable +param counts: +2,154,632,472 total +2,154,632,456 trainable +total_steps 16400 + +Done with model preparations! +param counts: +2,154,632,472 total +2,154,632,456 trainable +wandb mindeye run bold5k_v1 +wandb_config: + {'model_name': 'bold5k_v1', 'global_batch_size': '21', 'batch_size': 21, 'num_epochs': 80, 'num_sessions': 15, 'num_params': 2154632456, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': True, 'max_lr': 0.0003, 'mixup_pct': 0.33, 'num_samples_per_epoch': 4323, 'num_test': 480, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1} +wandb_id: bold5k_v1 +torch.Size([21, 1685]) torch.Size([21, 3, 224, 224]) torch.Size([21]) torch.Size([21]) +bold5k_v1 starting with epoch 0 / 80 + +===Finished!=== + diff --git a/MindEyeV2/src/slurms/545089.out b/MindEyeV2/src/slurms/545089.out new file mode 100644 index 0000000000000000000000000000000000000000..6a7c3030df107f7eba57d4f2f11de09f155f2329 --- /dev/null +++ b/MindEyeV2/src/slurms/545089.out @@ -0,0 +1,18 @@ +MASTER_ADDR=ip-10-0-133-32 +MASTER_PORT=12592 +WORLD_SIZE=1 +model_name=bold5k_nsdm1 +LOCAL RANK 0 +PID of this process = 840677 +device: cuda +Distributed environment: DistributedType.NO +Num processes: 1 +Process index: 0 +Local process index: 0 +Device: cuda + +Mixed precision type: fp16 + +distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16 +subj_list [1] num_sessions 15 +dividing batch size by subj_list, which will then be concatenated across subj during training... diff --git a/MindEyeV2/src/slurms/545090.err b/MindEyeV2/src/slurms/545090.err new file mode 100644 index 0000000000000000000000000000000000000000..d6bd498aa30ba3ab0b09768f10ccd78183b862d5 --- /dev/null +++ b/MindEyeV2/src/slurms/545090.err @@ -0,0 +1,14 @@ +[NbConvertApp] Converting notebook TrainB5k.ipynb to python +[NbConvertApp] Writing 52203 bytes to TrainB5k.py +wandb: Currently logged in as: ckadirt. Use `wandb login --relogin` to force relogin +wandb: wandb version 0.19.0 is available! To upgrade, please run: +wandb: $ pip install wandb --upgrade +wandb: Tracking run with wandb version 0.17.1 +wandb: Run data is saved locally in /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20241210_215527-bold5k_nsdm1 +wandb: Run `wandb offline` to turn off syncing. +wandb: Syncing run bold5k_nsdm1 +wandb: ⭐️ View project at https://stability.wandb.io/ckadirt/mindeye +wandb: 🚀 View run at https://stability.wandb.io/ckadirt/mindeye/runs/bold5k_nsdm1 + 0%| | 0/80 [00:00'} +2024-07-09 03:14:56,401 INFO MainThread:50989 [wandb_setup.py:_flush():76] Applying login settings: {'base_url': 'https://stability.wandb.io'} +2024-07-09 03:14:56,401 INFO MainThread:50989 [wandb_setup.py:_flush():76] Applying login settings: {} +2024-07-09 03:14:56,402 INFO MainThread:50989 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_031456-semantic_cluster_0.1/logs/debug.log +2024-07-09 03:14:56,405 INFO MainThread:50989 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_031456-semantic_cluster_0.1/logs/debug-internal.log +2024-07-09 03:14:56,406 INFO MainThread:50989 [wandb_init.py:_jupyter_setup():466] configuring jupyter hooks +2024-07-09 03:14:56,406 INFO MainThread:50989 [wandb_init.py:init():560] calling init triggers +2024-07-09 03:14:56,406 INFO MainThread:50989 [wandb_init.py:init():567] wandb.init called with sweep_config: {} +config: {'model_name': 'semantic_cluster_0.1', 'global_batch_size': 16, 'batch_size': 16, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 746793265, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 1e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'} +2024-07-09 03:14:56,406 INFO MainThread:50989 [wandb_init.py:init():610] starting backend +2024-07-09 03:14:56,406 INFO MainThread:50989 [wandb_init.py:init():614] setting up manager +2024-07-09 03:14:56,422 INFO MainThread:50989 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn +2024-07-09 03:14:56,429 INFO MainThread:50989 [wandb_init.py:init():622] backend started and connected +2024-07-09 03:14:56,454 INFO MainThread:50989 [wandb_run.py:_label_probe_notebook():1334] probe notebook +2024-07-09 03:14:56,456 INFO MainThread:50989 [wandb_run.py:_label_probe_notebook():1344] Unable to probe notebook: 'NoneType' object has no attribute 'get' +2024-07-09 03:14:56,456 INFO MainThread:50989 [wandb_init.py:init():711] updated telemetry +2024-07-09 03:14:56,568 INFO MainThread:50989 [wandb_init.py:init():744] communicating run to backend with 90.0 second timeout +2024-07-09 03:14:57,086 INFO MainThread:50989 [wandb_run.py:_on_init():2402] communicating current version +2024-07-09 03:14:57,140 INFO MainThread:50989 [wandb_run.py:_on_init():2411] got version response upgrade_message: "wandb version 0.17.4 is available! To upgrade, please run:\n $ pip install wandb --upgrade" + +2024-07-09 03:14:57,140 INFO MainThread:50989 [wandb_init.py:init():795] starting run threads in backend +2024-07-09 03:15:03,671 INFO MainThread:50989 [wandb_run.py:_console_start():2380] atexit reg +2024-07-09 03:15:03,671 INFO MainThread:50989 [wandb_run.py:_redirect():2235] redirect: wrap_raw +2024-07-09 03:15:03,672 INFO MainThread:50989 [wandb_run.py:_redirect():2300] Wrapping output streams. +2024-07-09 03:15:03,672 INFO MainThread:50989 [wandb_run.py:_redirect():2325] Redirects installed. +2024-07-09 03:15:03,679 INFO MainThread:50989 [wandb_init.py:init():838] run started, returning control to user process +2024-07-09 03:15:03,685 INFO MainThread:50989 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 03:15:03,686 INFO MainThread:50989 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 03:15:03,813 INFO MainThread:50989 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 03:15:03,814 INFO MainThread:50989 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 03:15:03,815 INFO MainThread:50989 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 03:15:03,892 INFO MainThread:50989 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 03:15:03,902 INFO MainThread:50989 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 03:15:03,903 INFO MainThread:50989 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 03:15:03,979 INFO MainThread:50989 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 03:15:04,261 INFO MainThread:50989 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 03:15:04,262 INFO MainThread:50989 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 03:15:04,349 INFO MainThread:50989 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 03:15:04,350 INFO MainThread:50989 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 03:15:04,351 INFO MainThread:50989 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 03:15:04,452 INFO MainThread:50989 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 03:22:21,534 INFO MainThread:50989 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 03:22:21,535 INFO MainThread:50989 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 03:23:58,503 INFO MainThread:50989 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 03:25:39,352 INFO MainThread:50989 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 03:25:39,354 INFO MainThread:50989 [wandb_init.py:_pause_backend():431] pausing backend diff --git a/MindEyeV2/src/wandb/run-20240709_031456-semantic_cluster_0.1/run-semantic_cluster_0.1.wandb b/MindEyeV2/src/wandb/run-20240709_031456-semantic_cluster_0.1/run-semantic_cluster_0.1.wandb new file mode 100644 index 0000000000000000000000000000000000000000..bf5d499f82bd99d74497c1a48bad4436b822394b Binary files /dev/null and b/MindEyeV2/src/wandb/run-20240709_031456-semantic_cluster_0.1/run-semantic_cluster_0.1.wandb differ diff --git a/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/code/_session_history.ipynb b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/code/_session_history.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..96db24afeea6d5848084bdab06c6b5a38f80175d --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/code/_session_history.ipynb @@ -0,0 +1,1036 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "770e88fe", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import sys\n", + "import json\n", + "import argparse\n", + "import numpy as np\n", + "import math\n", + "from einops import rearrange\n", + "import time\n", + "import random\n", + "import string\n", + "import h5py\n", + "from tqdm import tqdm\n", + "import webdataset as wds\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import torch\n", + "import torch.nn as nn\n", + "from torchvision import transforms\n", + "from accelerate import Accelerator\n", + "import torch.nn.functional as F\n", + "\n", + "# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main\n", + "sys.path.append('generative_models/')\n", + "import sgm\n", + "from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder\n", + "\n", + "# tf32 data type is faster than standard float32\n", + "torch.backends.cuda.matmul.allow_tf32 = True\n", + "\n", + "# custom functions #\n", + "import utils" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "38d13e72", + "metadata": {}, + "outputs": [], + "source": [ + "def classPrecision(logits, y_true, top=1):\n", + " \"\"\"\n", + " Calculate the precision of the top-n predictions.\n", + " \n", + " Parameters:\n", + " logits (torch.Tensor): The output logits from the model (shape: [batch_size, num_classes]).\n", + " y_true (torch.Tensor): The ground truth labels (shape: [batch_size]).\n", + " top (int): The number of top predictions to consider.\n", + " \n", + " Returns:\n", + " float: The precision percentage of the top-n predictions.\n", + " \"\"\"\n", + " # Apply softmax to get probabilities\n", + " probs = F.softmax(logits, dim=1).detach().cpu()\n", + " \n", + " # Get the top-n predictions\n", + " top_n_preds = torch.topk(probs, top, dim=1).indices.detach().cpu()\n", + "\n", + " # Move y_true to CPU and detach\n", + " y_true = y_true.detach().cpu()\n", + "\n", + " # Check if y_true is in top-n predictions\n", + " correct = top_n_preds.eq(y_true.view(-1, 1).expand_as(top_n_preds))\n", + "\n", + " # Calculate precision\n", + " precision = correct.sum().item() / y_true.size(0)\n", + " \n", + " return precision * 100\n", + "\n", + "# Example usage:\n", + "logits = torch.randn(8, 41) # Example logits tensor\n", + "y_true = torch.randint(0, 41, (8,)) # Example ground truth labels\n", + "\n", + "top_n_precision = classPrecision(logits, y_true, top=1)\n", + "print(f\"Top-1 Precision: {top_n_precision:.2f}%\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f8410934", + "metadata": {}, + "outputs": [], + "source": [ + "### Multi-GPU config ###\n", + "local_rank = os.getenv('RANK')\n", + "if local_rank is None: \n", + " local_rank = 0\n", + "else:\n", + " local_rank = int(local_rank)\n", + "print(\"LOCAL RANK \", local_rank) \n", + "\n", + "data_type = torch.float16 # change depending on your mixed_precision\n", + "num_devices = torch.cuda.device_count()\n", + "if num_devices==0: num_devices = 1\n", + "\n", + "# First use \"accelerate config\" in terminal and setup using deepspeed stage 2 with CPU offloading!\n", + "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n", + "if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job\n", + " global_batch_size = batch_size = 16\n", + "else:\n", + " global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n", + " batch_size = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a02d706c", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"PID of this process =\",os.getpid())\n", + "device = accelerator.device\n", + "print(\"device:\",device)\n", + "world_size = accelerator.state.num_processes\n", + "distributed = not accelerator.state.distributed_type == 'NO'\n", + "num_devices = torch.cuda.device_count()\n", + "if num_devices==0 or not distributed: num_devices = 1\n", + "num_workers = num_devices\n", + "print(accelerator.state)\n", + "\n", + "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\n", + "print = accelerator.print # only print if local_rank=0" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ab031a50", + "metadata": {}, + "outputs": [], + "source": [ + "# if running this interactively, can specify jupyter_args here for argparser to use\n", + "if utils.is_interactive():\n", + " model_name = \"semantic_cluster_0.1\"\n", + " print(\"model_name:\", model_name)\n", + " \n", + " # global_batch_size and batch_size should already be defined in the 2nd cell block\n", + " jupyter_args = f\"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \\\n", + " --cache_dir=/weka/proj-medarc/shared/cache \\\n", + " --model_name={model_name} \\\n", + " --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \\\n", + " --hidden_dim=1024 --clip_scale=1. \\\n", + " --no-blurry_recon --blur_scale=.5 \\\n", + " --use_prior --prior_scale=30 \\\n", + " --n_blocks=4 --max_lr=1e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \\\n", + " --ckpt_interval=999 --no-ckpt_saving --wandb_log\"\n", + " # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow\n", + "\n", + " print(jupyter_args)\n", + " jupyter_args = jupyter_args.split()\n", + " \n", + " from IPython.display import clear_output # function to clear print outputs in cell\n", + " %load_ext autoreload \n", + " # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions\n", + " %autoreload 2 " + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "afbef21c", + "metadata": {}, + "outputs": [], + "source": [ + "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n", + "parser.add_argument(\n", + " \"--model_name\", type=str, default=\"testing2\",\n", + " help=\"name of model, used for ckpt saving and wandb logging (if enabled)\",\n", + ")\n", + "parser.add_argument(\n", + " \"--data_path\", type=str, default=os.getcwd(),\n", + " help=\"Path to where NSD data is stored / where to download it to\",\n", + ")\n", + "parser.add_argument(\n", + " \"--cache_dir\", type=str, default=os.getcwd(),\n", + " help=\"Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.\",\n", + ")\n", + "parser.add_argument(\n", + " \"--subj\",type=int, default=1, choices=[1,2,3,4,5,6,7,8],\n", + " help=\"Validate on which subject?\",\n", + ")\n", + "parser.add_argument(\n", + " \"--multisubject_ckpt\", type=str, default=None,\n", + " help=\"Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.\",\n", + ")\n", + "parser.add_argument(\n", + " \"--num_sessions\", type=int, default=1,\n", + " help=\"Number of training sessions to include\",\n", + ")\n", + "parser.add_argument(\n", + " \"--use_prior\",action=argparse.BooleanOptionalAction,default=True,\n", + " help=\"whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)\",\n", + ")\n", + "parser.add_argument(\n", + " \"--batch_size\", type=int, default=16,\n", + " help=\"Batch size can be increased by 10x if only training retreival submodule and not diffusion prior\",\n", + ")\n", + "parser.add_argument(\n", + " \"--wandb_log\",action=argparse.BooleanOptionalAction,default=False,\n", + " help=\"whether to log to wandb\",\n", + ")\n", + "parser.add_argument(\n", + " \"--wandb_project\",type=str,default=\"stability\",\n", + " help=\"wandb project name\",\n", + ")\n", + "parser.add_argument(\n", + " \"--mixup_pct\",type=float,default=.33,\n", + " help=\"proportion of way through training when to switch from BiMixCo to SoftCLIP\",\n", + ")\n", + "parser.add_argument(\n", + " \"--blurry_recon\",action=argparse.BooleanOptionalAction,default=True,\n", + " help=\"whether to output blurry reconstructions\",\n", + ")\n", + "parser.add_argument(\n", + " \"--blur_scale\",type=float,default=.5,\n", + " help=\"multiply loss from blurry recons by this number\",\n", + ")\n", + "parser.add_argument(\n", + " \"--clip_scale\",type=float,default=1.,\n", + " help=\"multiply contrastive loss by this number\",\n", + ")\n", + "parser.add_argument(\n", + " \"--prior_scale\",type=float,default=30,\n", + " help=\"multiply diffusion prior loss by this\",\n", + ")\n", + "parser.add_argument(\n", + " \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=False,\n", + " help=\"whether to use image augmentation\",\n", + ")\n", + "parser.add_argument(\n", + " \"--num_epochs\",type=int,default=150,\n", + " help=\"number of epochs of training\",\n", + ")\n", + "parser.add_argument(\n", + " \"--multi_subject\",action=argparse.BooleanOptionalAction,default=False,\n", + ")\n", + "parser.add_argument(\n", + " \"--new_test\",action=argparse.BooleanOptionalAction,default=True,\n", + ")\n", + "parser.add_argument(\n", + " \"--n_blocks\",type=int,default=4,\n", + ")\n", + "parser.add_argument(\n", + " \"--hidden_dim\",type=int,default=1024,\n", + ")\n", + "parser.add_argument(\n", + " \"--lr_scheduler_type\",type=str,default='cycle',choices=['cycle','linear'],\n", + ")\n", + "parser.add_argument(\n", + " \"--ckpt_saving\",action=argparse.BooleanOptionalAction,default=True,\n", + ")\n", + "parser.add_argument(\n", + " \"--ckpt_interval\",type=int,default=5,\n", + " help=\"save backup ckpt and reconstruct every x epochs\",\n", + ")\n", + "parser.add_argument(\n", + " \"--seed\",type=int,default=42,\n", + ")\n", + "parser.add_argument(\n", + " \"--max_lr\",type=float,default=3e-5,\n", + ")\n", + "\n", + "if utils.is_interactive():\n", + " args = parser.parse_args(jupyter_args)\n", + "else:\n", + " args = parser.parse_args()\n", + "\n", + "# create global variables without the args prefix\n", + "for attribute_name in vars(args).keys():\n", + " globals()[attribute_name] = getattr(args, attribute_name)\n", + " \n", + "# seed all random functions\n", + "utils.seed_everything(seed)\n", + "\n", + "outdir = os.path.abspath(f'../train_logs/{model_name}')\n", + "if not os.path.exists(outdir) and ckpt_saving:\n", + " os.makedirs(outdir,exist_ok=True)\n", + " \n", + "if use_image_aug or blurry_recon:\n", + " import kornia\n", + " from kornia.augmentation.container import AugmentationSequential\n", + "if use_image_aug:\n", + " img_augment = AugmentationSequential(\n", + " kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),\n", + " same_on_batch=False,\n", + " data_keys=[\"input\"],\n", + " )\n", + " \n", + "if multi_subject:\n", + " subj_list = np.arange(1,9)\n", + " subj_list = subj_list[subj_list != subj]\n", + "else:\n", + " subj_list = [subj]\n", + "\n", + "print(\"subj_list\", subj_list, \"num_sessions\", num_sessions)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "636f61c6", + "metadata": {}, + "outputs": [], + "source": [ + "def my_split_by_node(urls): return urls\n", + "num_voxels_list = []\n", + "\n", + "if multi_subject:\n", + " nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])\n", + " num_samples_per_epoch = (750*40) // num_devices \n", + "else:\n", + " num_samples_per_epoch = (750*num_sessions) // num_devices \n", + "\n", + "print(\"dividing batch size by subj_list, which will then be concatenated across subj during training...\") \n", + "batch_size = batch_size // len(subj_list)\n", + "\n", + "num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))\n", + "\n", + "print(\"batch_size =\", batch_size, \"num_iterations_per_epoch =\",num_iterations_per_epoch, \"num_samples_per_epoch =\",num_samples_per_epoch)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "55dcf898", + "metadata": {}, + "outputs": [], + "source": [ + "train_data = {}\n", + "train_dl = {}\n", + "num_voxels = {}\n", + "voxels = {}\n", + "for s in subj_list:\n", + " print(f\"Training with {num_sessions} sessions\")\n", + " if multi_subject:\n", + " train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{nsessions_allsubj[s-1]-1}\" + \"}.tar\"\n", + " else:\n", + " train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n", + " print(train_url)\n", + " \n", + " train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,nodesplitter=my_split_by_node)\\\n", + " .shuffle(750, initial=1500, rng=random.Random(42))\\\n", + " .decode(\"torch\")\\\n", + " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", + " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", + " train_dl[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n", + "\n", + " f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n", + " betas = f['betas'][:]\n", + " betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n", + " num_voxels_list.append(betas[0].shape[-1])\n", + " num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n", + " voxels[f'subj0{s}'] = betas\n", + " print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n", + "\n", + "print(\"Loaded all subj train dls and betas!\\n\")\n", + "\n", + "# Validate only on one subject\n", + "if multi_subject: \n", + " subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list\n", + "if not new_test: # using old test set from before full dataset released (used in original MindEye paper)\n", + " if subj==3:\n", + " num_test=2113\n", + " elif subj==4:\n", + " num_test=1985\n", + " elif subj==6:\n", + " num_test=2113\n", + " elif subj==8:\n", + " num_test=1985\n", + " else:\n", + " num_test=2770\n", + " test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n", + "elif new_test: # using larger test set from after full dataset released\n", + " if subj==3:\n", + " num_test=2371\n", + " elif subj==4:\n", + " num_test=2188\n", + " elif subj==6:\n", + " num_test=2371\n", + " elif subj==8:\n", + " num_test=2188\n", + " else:\n", + " num_test=3000\n", + " test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n", + "print(test_url)\n", + "test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\\\n", + " .shuffle(750, initial=1500, rng=random.Random(42))\\\n", + " .decode(\"torch\")\\\n", + " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", + " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", + "test_dl = torch.utils.data.DataLoader(test_data, batch_size=num_test, shuffle=False, drop_last=True, pin_memory=True)\n", + "print(f\"Loaded test dl for subj{subj}!\\n\")" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "3ac1391b", + "metadata": {}, + "outputs": [], + "source": [ + "# Load 73k NSD images\n", + "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n", + "images = f['images']\n", + "print(\"Loaded all 73k possible NSD images to cpu!\", images.shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "4f4acf44", + "metadata": {}, + "outputs": [], + "source": [ + "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n", + " arch=\"ViT-bigG-14\",\n", + " version=\"laion2b_s39b_b160k\",\n", + " output_tokens=True,\n", + " only_tokens=True,\n", + ")\n", + "clip_img_embedder.to(device)\n", + "\n", + "clip_seq_dim = 256\n", + "clip_emb_dim = 1664" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "9bc03c5e", + "metadata": {}, + "outputs": [], + "source": [ + "if blurry_recon:\n", + " from diffusers import AutoencoderKL \n", + " autoenc = AutoencoderKL(\n", + " down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],\n", + " up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],\n", + " block_out_channels=[128, 256, 512, 512],\n", + " layers_per_block=2,\n", + " sample_size=256,\n", + " )\n", + " ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')\n", + " autoenc.load_state_dict(ckpt)\n", + " \n", + " autoenc.eval()\n", + " autoenc.requires_grad_(False)\n", + " autoenc.to(device)\n", + " utils.count_params(autoenc)\n", + " \n", + " from autoencoder.convnext import ConvnextXL\n", + " cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')\n", + " cnx.requires_grad_(False)\n", + " cnx.eval()\n", + " cnx.to(device)\n", + " \n", + " mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)\n", + " std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)\n", + " \n", + " blur_augs = AugmentationSequential(\n", + " kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),\n", + " kornia.augmentation.RandomGrayscale(p=0.1),\n", + " kornia.augmentation.RandomSolarize(p=0.1),\n", + " kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),\n", + " data_keys=[\"input\"],\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "fbf1e6fb", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "MindEyeModule()" + ] + } + ], + "source": [ + "class MindEyeModule(nn.Module):\n", + " def __init__(self):\n", + " super(MindEyeModule, self).__init__()\n", + " def forward(self, x):\n", + " return x\n", + " \n", + "model = MindEyeModule()\n", + "model" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "f291a3f5", + "metadata": {}, + "outputs": [], + "source": [ + "class RidgeRegression(torch.nn.Module):\n", + " # make sure to add weight_decay when initializing optimizer to enable regularization\n", + " def __init__(self, input_sizes, out_features): \n", + " super(RidgeRegression, self).__init__()\n", + " self.out_features = out_features\n", + " self.linears = torch.nn.ModuleList([\n", + " torch.nn.Linear(input_size, out_features) for input_size in input_sizes\n", + " ])\n", + " def forward(self, x, subj_idx):\n", + " out = self.linears[subj_idx](x[:,0]).unsqueeze(1)\n", + " return out\n", + " \n", + "class IndividRidgeRegression(torch.nn.Module):\n", + " def __init__(self, input_size, out_features):\n", + " super(IndividRidgeRegression, self).__init__()\n", + " self.out_features = out_features\n", + " self.linear = torch.nn.Linear(input_size, out_features)\n", + " def forward(self, x):\n", + " out = self.linear(x)\n", + " return out\n", + " \n", + "model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)\n", + "utils.count_params(model.ridge)\n", + "utils.count_params(model)\n", + "\n", + "# test on subject 1 with fake data\n", + "b = torch.randn((2,1,num_voxels_list[0]))\n", + "print(b.shape, model.ridge(b,0).shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "1eec72cb", + "metadata": {}, + "outputs": [], + "source": [ + "from models import BrainNetwork\n", + "model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,\n", + " clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, \n", + " blurry_recon=blurry_recon, clip_scale=clip_scale)\n", + "utils.count_params(model.backbone)\n", + "utils.count_params(model)\n", + "\n", + "# test that the model works on some fake data\n", + "b = torch.randn((2,1,hidden_dim))\n", + "print(\"b.shape\",b.shape)\n", + "\n", + "backbone_, clip_, blur_ = model.backbone(b)\n", + "print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "f72a54e8", + "metadata": {}, + "outputs": [], + "source": [ + "if use_prior:\n", + " from models import *\n", + "\n", + " # setup diffusion prior network\n", + " out_dim = clip_emb_dim\n", + " depth = 6\n", + " dim_head = 52\n", + " heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim\n", + " timesteps = 100\n", + "\n", + " prior_network = PriorNetwork(\n", + " dim=out_dim,\n", + " depth=depth,\n", + " dim_head=dim_head,\n", + " heads=heads,\n", + " causal=False,\n", + " num_tokens = clip_seq_dim,\n", + " learned_query_mode=\"pos_emb\"\n", + " )\n", + "\n", + " model.diffusion_prior = BrainDiffusionPrior(\n", + " net=prior_network,\n", + " image_embed_dim=out_dim,\n", + " condition_on_text_encodings=False,\n", + " timesteps=timesteps,\n", + " cond_drop_prob=0.2,\n", + " image_embed_scale=None,\n", + " )\n", + " \n", + " utils.count_params(model.diffusion_prior)\n", + " utils.count_params(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "57803a48", + "metadata": {}, + "outputs": [], + "source": [ + "path_semantic_names = \"/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy\"\n", + "path_semantic_cluster = \"/weka/proj-fmri/ckadirt/MindEyeV2/src/COCO_73k_semantic_cluster.npy\"\n", + "semantic_cluster_names = np.load(path_semantic_names)\n", + "semantic_cluster = np.load(path_semantic_cluster)\n", + "possible_semantic_clusters = np.unique(semantic_cluster)\n", + "\n", + "# one-hot encode semantic clusters\n", + "# move possible_semantic_clusters to numbers and create a dictionary\n", + "semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)}\n", + "semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters)))\n", + "for i, cluster in enumerate(semantic_cluster):\n", + " semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1\n", + "\n", + "\n", + "print(\"semantic_cluster_onehot.shape\", semantic_cluster_onehot.shape)\n", + "\n", + "num_seman_clusters = len(np.unique(semantic_cluster))\n", + "print(\"num_seman_clusters\", num_seman_clusters)" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "4497c76c", + "metadata": {}, + "outputs": [], + "source": [ + "# plot some images next to their semantic cluster\n", + "fig, ax = plt.subplots(1, 5, figsize=(20, 4))\n", + "for i in range(5):\n", + " # covert numpy array images to float32 \n", + " image_index = torch.randint(0, len(images), (1,)).item()\n", + " print(image_index)\n", + " ax[i].imshow(images[image_index].transpose(1,2,0).astype(np.float32))\n", + " ax[i].set_title(semantic_cluster[image_index])\n", + " ax[i].axis(\"off\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "754f285d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "746793265" + ] + } + ], + "source": [ + "model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters)\n", + "utils.count_params(model.RRClassifier)\n", + "utils.count_params(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "7be4a60e", + "metadata": {}, + "outputs": [], + "source": [ + "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n", + "\n", + "opt_grouped_parameters = [\n", + " {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},\n", + " {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n", + " {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n", + " {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1e-2},\n", + "]\n", + "# if use_prior:\n", + "# opt_grouped_parameters.extend([\n", + "# {'params': [p for n, p in model.diffusion_prior.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n", + "# {'params': [p for n, p in model.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n", + "# ])\n", + "# opt_grouped_parameters.extend([\n", + "# \n", + "# ])\n", + "\n", + "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)\n", + "\n", + "if lr_scheduler_type == 'linear':\n", + " lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n", + " optimizer,\n", + " total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),\n", + " last_epoch=-1\n", + " )\n", + "elif lr_scheduler_type == 'cycle':\n", + " total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))\n", + " print(\"total_steps\", total_steps)\n", + " lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n", + " optimizer, \n", + " max_lr=max_lr,\n", + " total_steps=total_steps,\n", + " final_div_factor=1000,\n", + " last_epoch=-1, pct_start=2/num_epochs\n", + " )\n", + " \n", + "def save_ckpt(tag):\n", + " ckpt_path = outdir+f'/{tag}.pth'\n", + " if accelerator.is_main_process:\n", + " unwrapped_model = accelerator.unwrap_model(model)\n", + " torch.save({\n", + " 'epoch': epoch,\n", + " 'model_state_dict': unwrapped_model.state_dict(),\n", + " 'optimizer_state_dict': optimizer.state_dict(),\n", + " 'lr_scheduler': lr_scheduler.state_dict(),\n", + " 'train_losses': losses,\n", + " 'test_losses': test_losses,\n", + " 'lrs': lrs,\n", + " }, ckpt_path)\n", + " print(f\"\\n---saved {outdir}/{tag} ckpt!---\\n\")\n", + "\n", + "def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): \n", + " print(f\"\\n---loading {outdir}/{tag}.pth ckpt---\\n\")\n", + " checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n", + " state_dict = checkpoint['model_state_dict']\n", + " if multisubj_loading: # remove incompatible ridge layer that will otherwise error\n", + " state_dict.pop('ridge.linears.0.weight',None)\n", + " model.load_state_dict(state_dict, strict=strict)\n", + " if load_epoch:\n", + " globals()[\"epoch\"] = checkpoint['epoch']\n", + " print(\"Epoch\",epoch)\n", + " if load_optimizer:\n", + " optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n", + " if load_lr:\n", + " lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n", + " del checkpoint\n", + "\n", + "print(\"\\nDone with model preparations!\")\n", + "num_params = utils.count_params(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "6fedf0ad", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1e-05" + ] + } + ], + "source": [ + "max_lr" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "7d78e7f3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "wandb version 0.17.4 is available! To upgrade, please run:\n", + " $ pip install wandb --upgrade" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Tracking run with wandb version 0.17.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Run data is saved locally in /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Resuming run semantic_cluster_0.1 to Weights & Biases (docs)
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + " View project at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster_0.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + " View run at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster_0.1/runs/semantic_cluster_0.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "if local_rank==0 and wandb_log: # only use main process for wandb logging\n", + " import wandb\n", + " wandb_project = 'mindeye_semantic_cluster_0.1'\n", + " print(f\"wandb {wandb_project} run {model_name}\")\n", + " # need to configure wandb beforehand in terminal with \"wandb init\"!\n", + " wandb_config = {\n", + " \"model_name\": model_name,\n", + " \"global_batch_size\": global_batch_size,\n", + " \"batch_size\": batch_size,\n", + " \"num_epochs\": num_epochs,\n", + " \"num_sessions\": num_sessions,\n", + " \"num_params\": num_params,\n", + " \"clip_scale\": clip_scale,\n", + " \"prior_scale\": prior_scale,\n", + " \"blur_scale\": blur_scale,\n", + " \"use_image_aug\": use_image_aug,\n", + " \"max_lr\": max_lr,\n", + " \"mixup_pct\": mixup_pct,\n", + " \"num_samples_per_epoch\": num_samples_per_epoch,\n", + " \"num_test\": num_test,\n", + " \"ckpt_interval\": ckpt_interval,\n", + " \"ckpt_saving\": ckpt_saving,\n", + " \"seed\": seed,\n", + " \"distributed\": distributed,\n", + " \"num_devices\": num_devices,\n", + " \"world_size\": world_size,\n", + " \"train_url\": train_url,\n", + " \"test_url\": test_url,\n", + " }\n", + " print(\"wandb_config:\\n\",wandb_config)\n", + " print(\"wandb_id:\",model_name)\n", + " wandb.login(host='https://stability.wandb.io')\n", + " wandb.init(\n", + " id=model_name,\n", + " project=wandb_project,\n", + " name=model_name,\n", + " config=wandb_config,\n", + " resume=\"allow\",\n", + " )\n", + "else:\n", + " wandb_log = False" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "69011870", + "metadata": {}, + "outputs": [], + "source": [ + "epoch = 0\n", + "losses, test_losses, lrs = [], [], []\n", + "best_test_loss = 1e9\n", + "torch.cuda.empty_cache()" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "cbc20574", + "metadata": {}, + "outputs": [], + "source": [ + "# load multisubject stage1 ckpt if set\n", + "if multisubject_ckpt is not None:\n", + " load_ckpt(\"last\",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "d89584ca", + "metadata": {}, + "outputs": [], + "source": [ + "train_dls = [train_dl[f'subj0{s}'] for s in subj_list]\n", + "\n", + "model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot)\n", + "# leaving out test_dl since we will only have local_rank 0 device do evals" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "47b1e63b", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_semantic_clusters(images, indexes, semantic_cluster, semantic_cluster_dict):\n", + " # check images and indexes are the same length\n", + " assert len(images) == len(indexes)\n", + " \n", + " # if images are tensors, convert them to numpy arrays and move them to cpu\n", + " if isinstance(images, torch.Tensor):\n", + " images = images.cpu().numpy()\n", + "\n", + " fig, ax = plt.subplots(1, len(images), figsize=(20, 4))\n", + " for i, index in enumerate(indexes):\n", + " # covert numpy array images to float32 \n", + " ax[i].imshow(images[i].transpose(1,2,0).astype(np.float32))\n", + " # search the key in the dictionary based on the index as value\n", + " name = {i for i in semantic_cluster_dict if semantic_cluster_dict[i] == index}\n", + " ax[i].set_title(name)\n", + " ax[i].axis(\"off\")\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "fd22d30b", + "metadata": {}, + "outputs": [], + "source": [ + "RTTT" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "d351dffe", + "metadata": {}, + "outputs": [], + "source": [ + "if local_rank==0 and wandb_log: # only use main process for wandb logging\n", + " import wandb\n", + " wandb_project = 'mindeye_semantic_cluster_0.2'\n", + " print(f\"wandb {wandb_project} run {model_name}\")\n", + " # need to configure wandb beforehand in terminal with \"wandb init\"!\n", + " wandb_config = {\n", + " \"model_name\": model_name,\n", + " \"global_batch_size\": global_batch_size,\n", + " \"batch_size\": batch_size,\n", + " \"num_epochs\": num_epochs,\n", + " \"num_sessions\": num_sessions,\n", + " \"num_params\": num_params,\n", + " \"clip_scale\": clip_scale,\n", + " \"prior_scale\": prior_scale,\n", + " \"blur_scale\": blur_scale,\n", + " \"use_image_aug\": use_image_aug,\n", + " \"max_lr\": max_lr,\n", + " \"mixup_pct\": mixup_pct,\n", + " \"num_samples_per_epoch\": num_samples_per_epoch,\n", + " \"num_test\": num_test,\n", + " \"ckpt_interval\": ckpt_interval,\n", + " \"ckpt_saving\": ckpt_saving,\n", + " \"seed\": seed,\n", + " \"distributed\": distributed,\n", + " \"num_devices\": num_devices,\n", + " \"world_size\": world_size,\n", + " \"train_url\": train_url,\n", + " \"test_url\": test_url,\n", + " }\n", + " print(\"wandb_config:\\n\",wandb_config)\n", + " print(\"wandb_id:\",model_name)\n", + " wandb.login(host='https://stability.wandb.io')\n", + " wandb.init(\n", + " id=model_name,\n", + " project=wandb_project,\n", + " name=model_name,\n", + " config=wandb_config,\n", + " resume=\"allow\",\n", + " )\n", + "else:\n", + " wandb_log = False" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "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.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/config.yaml b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1c5c95828b87b66cdff684dd99f6d6afe751354d --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/config.yaml @@ -0,0 +1,119 @@ +wandb_version: 1 + +model_name: + desc: null + value: semantic_cluster_0.1 +global_batch_size: + desc: null + value: 16 +batch_size: + desc: null + value: 16 +num_epochs: + desc: null + value: 150 +num_sessions: + desc: null + value: 40 +num_params: + desc: null + value: 746793265 +clip_scale: + desc: null + value: 1.0 +prior_scale: + desc: null + value: 30.0 +blur_scale: + desc: null + value: 0.5 +use_image_aug: + desc: null + value: false +max_lr: + desc: null + value: 1.0e-05 +mixup_pct: + desc: null + value: 0.33 +num_samples_per_epoch: + desc: null + value: 30000 +num_test: + desc: null + value: 3000 +ckpt_interval: + desc: null + value: 999 +ckpt_saving: + desc: null + value: false +seed: + desc: null + value: 42 +distributed: + desc: null + value: false +num_devices: + desc: null + value: 1 +world_size: + desc: null + value: 1 +train_url: + desc: null + value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar +test_url: + desc: null + value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar +_wandb: + desc: null + value: + python_version: 3.11.9 + cli_version: 0.17.1 + framework: huggingface + huggingface_version: 4.37.2 + is_jupyter_run: true + is_kaggle_kernel: false + start_time: 1720532363 + t: + 1: + - 1 + - 9 + - 11 + - 41 + - 49 + - 55 + - 63 + - 71 + - 79 + - 83 + - 103 + 2: + - 1 + - 9 + - 11 + - 41 + - 49 + - 55 + - 63 + - 71 + - 79 + - 83 + - 103 + 3: + - 2 + - 5 + - 13 + - 14 + - 16 + - 23 + - 62 + 4: 3.11.9 + 5: 0.17.1 + 6: 4.37.2 + 8: + - 1 + - 5 + 13: linux-x86_64 + session_history: code/_session_history.ipynb diff --git a/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..3161ad60260d4d7f13ee04255338e3994721bd2c --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log @@ -0,0 +1,6 @@ + +wandb: WARNING Calling wandb.login() after wandb.init() has no effect. +wandb mindeye_semantic_cluster_0.2 run semantic_cluster_0.1 +wandb_config: + {'model_name': 'semantic_cluster_0.1', 'global_batch_size': 16, 'batch_size': 16, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 746793265, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 1e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'} +wandb_id: semantic_cluster_0.1 \ No newline at end of file diff --git a/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..0bd25aa739892a3186231ebaab1b00dd803e51d9 --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json @@ -0,0 +1 @@ +{"train/loss_clip_total": 0, "test/class_precisions_5": 16, "train/class_precisions_1": 0, "_runtime": 624.9114198684692, "train/loss": 50.06567057495117, "train/loss_RR": 0, "test/recon_mse": 0, "test/loss_prior": 1.385520577430725, "train/num_steps": 1875, "train/blurry_pixcorr": 0, "train/loss_blurry_total": 0, "train/class_precisions_5": 0, "test/recon_cossim": 0, "test/loss_clip_total": 4.698548316955566, "train/recon_mse": 0, "train/fwd_pct_correct": 0, "train/loss_blurry_cont_total": 0, "_step": 0, "test/loss": 50.87119674682617, "test/loss_RR": 4.607032775878906, "test/class_precisions_1": 4.333333333333334, "train/lr": 5.181899642439222e-06, "test/num_steps": 1, "train/loss_prior": 0, "test/test_bwd_pct_correct": 0.06000000238418579, "train/class_precisions_10": 0, "_timestamp": 1720495521.3412728, "train/recon_cossim": 0, "test/blurry_pixcorr": 0, "train/bwd_pct_correct": 0, "test/class_precisions_10": 28.000000000000004, "test/test_fwd_pct_correct": 0.12333333492279051, "_wandb": {"runtime": 6}} \ No newline at end of file diff --git a/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug-internal.log b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..48f4c91e040ff549f487992f98aa182dc57ffaa4 --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug-internal.log @@ -0,0 +1,303 @@ +2024-07-09 13:39:23,520 INFO StreamThr :1763125 [internal.py:wandb_internal():85] W&B internal server running at pid: 1763125, started at: 2024-07-09 13:39:23.511074 +2024-07-09 13:39:23,522 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status +2024-07-09 13:39:23,542 INFO WriterThread:1763125 [datastore.py:open_for_write():87] open: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/run-semantic_cluster_0.1.wandb +2024-07-09 13:39:23,548 DEBUG SenderThread:1763125 [sender.py:send():379] send: header +2024-07-09 13:39:23,660 DEBUG SenderThread:1763125 [sender.py:send():379] send: run +2024-07-09 13:39:23,674 INFO SenderThread:1763125 [sender.py:_setup_resume():749] checking resume status for None/mindeye_semantic_cluster_0.1/semantic_cluster_0.1 +2024-07-09 13:39:23,990 INFO SenderThread:1763125 [sender.py:_setup_resume():829] configured resuming with: ResumeState(resumed=True,step=1,history=1,events=22,output=69,runtime=642.926949,wandb_runtime=None,summary={'train/loss_clip_total': 0, 'test/class_precisions_5': 16, 'train/class_precisions_1': 0, '_runtime': 624.9114198684692, 'train/loss': 50.06567057495117, 'train/loss_RR': 0, 'test/recon_mse': 0, 'test/loss_prior': 1.385520577430725, 'train/num_steps': 1875, 'train/blurry_pixcorr': 0, 'train/loss_blurry_total': 0, 'train/class_precisions_5': 0, 'test/recon_cossim': 0, 'test/loss_clip_total': 4.698548316955566, 'train/recon_mse': 0, 'train/fwd_pct_correct': 0, 'train/loss_blurry_cont_total': 0, '_step': 0, 'test/loss': 50.87119674682617, 'test/loss_RR': 4.607032775878906, 'test/class_precisions_1': 4.333333333333334, 'train/lr': 5.181899642439222e-06, 'test/num_steps': 1, 'train/loss_prior': 0, 'test/test_bwd_pct_correct': 0.06000000238418579, 'train/class_precisions_10': 0, '_timestamp': 1720495521.3412728, 'train/recon_cossim': 0, 'test/blurry_pixcorr': 0, 'train/bwd_pct_correct': 0, 'test/class_precisions_10': 28.000000000000004, 'test/test_fwd_pct_correct': 0.12333333492279051},config={'seed': {'desc': None, 'value': 42}, '_wandb': {'desc': None, 'value': {'t': {'1': [1, 9, 11, 41, 49, 55, 63, 71, 79, 83, 103], '2': [1, 9, 11, 41, 49, 55, 63, 71, 79, 83, 103], '3': [13, 14, 16, 23], '4': '3.11.9', '5': '0.17.1', '6': '4.37.2', '8': [1, 5], '13': 'linux-x86_64'}, 'framework': 'huggingface', 'start_time': 1720494896, 'cli_version': '0.17.1', 'is_jupyter_run': True, 'python_version': '3.11.9', 'is_kaggle_kernel': False, 'huggingface_version': '4.37.2'}}, 'max_lr': {'desc': None, 'value': 1e-05}, 'num_test': {'desc': None, 'value': 3000}, 'test_url': {'desc': None, 'value': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'}, 'mixup_pct': {'desc': None, 'value': 0.33}, 'train_url': {'desc': None, 'value': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar'}, 'batch_size': {'desc': None, 'value': 16}, 'blur_scale': {'desc': None, 'value': 0.5}, 'clip_scale': {'desc': None, 'value': 1}, 'model_name': {'desc': None, 'value': 'semantic_cluster_0.1'}, 'num_epochs': {'desc': None, 'value': 150}, 'num_params': {'desc': None, 'value': 746793265}, 'world_size': {'desc': None, 'value': 1}, 'ckpt_saving': {'desc': None, 'value': False}, 'distributed': {'desc': None, 'value': False}, 'num_devices': {'desc': None, 'value': 1}, 'prior_scale': {'desc': None, 'value': 30}, 'num_sessions': {'desc': None, 'value': 40}, 'ckpt_interval': {'desc': None, 'value': 999}, 'use_image_aug': {'desc': None, 'value': False}, 'global_batch_size': {'desc': None, 'value': 16}, 'num_samples_per_epoch': {'desc': None, 'value': 30000}},tags=[]) +2024-07-09 13:39:24,197 INFO SenderThread:1763125 [dir_watcher.py:__init__():211] watching files in: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files +2024-07-09 13:39:24,197 INFO SenderThread:1763125 [sender.py:_start_run_threads():1188] run started: semantic_cluster_0.1 with start time 1720531720.5911 +2024-07-09 13:39:24,197 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: summary_record +2024-07-09 13:39:24,208 INFO SenderThread:1763125 [sender.py:_save_file():1454] saving file wandb-summary.json with policy end +2024-07-09 13:39:24,219 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: check_version +2024-07-09 13:39:24,220 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: check_version +2024-07-09 13:39:24,290 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: run_start +2024-07-09 13:39:24,486 DEBUG HandlerThread:1763125 [system_info.py:__init__():26] System info init +2024-07-09 13:39:24,486 DEBUG HandlerThread:1763125 [system_info.py:__init__():41] System info init done +2024-07-09 13:39:24,486 INFO HandlerThread:1763125 [system_monitor.py:start():194] Starting system monitor +2024-07-09 13:39:24,486 INFO SystemMonitor:1763125 [system_monitor.py:_start():158] Starting system asset monitoring threads +2024-07-09 13:39:24,487 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started cpu monitoring +2024-07-09 13:39:24,487 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started disk monitoring +2024-07-09 13:39:24,488 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started gpu monitoring +2024-07-09 13:39:24,488 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started memory monitoring +2024-07-09 13:39:24,489 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started network monitoring +2024-07-09 13:39:25,117 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: python_packages +2024-07-09 13:39:25,118 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: python_packages +2024-07-09 13:39:25,118 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: stop_status +2024-07-09 13:39:25,119 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: internal_messages +2024-07-09 13:39:25,132 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: pause +2024-07-09 13:39:25,133 INFO HandlerThread:1763125 [handler.py:handle_request_pause():724] stopping system metrics thread +2024-07-09 13:39:25,133 INFO HandlerThread:1763125 [system_monitor.py:finish():203] Stopping system monitor +2024-07-09 13:39:25,133 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():172] Starting system metrics aggregation loop +2024-07-09 13:39:25,134 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():179] Finished system metrics aggregation loop +2024-07-09 13:39:25,134 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():183] Publishing last batch of metrics +2024-07-09 13:39:25,135 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined cpu monitor +2024-07-09 13:39:25,135 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined disk monitor +2024-07-09 13:39:25,136 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: stop_status +2024-07-09 13:39:25,204 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/requirements.txt +2024-07-09 13:39:25,204 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json +2024-07-09 13:39:25,316 DEBUG SenderThread:1763125 [sender.py:send():379] send: telemetry +2024-07-09 13:39:26,945 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined gpu monitor +2024-07-09 13:39:26,945 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined memory monitor +2024-07-09 13:39:26,945 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined network monitor +2024-07-09 13:39:26,945 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: resume +2024-07-09 13:39:26,945 INFO HandlerThread:1763125 [handler.py:handle_request_resume():715] starting system metrics thread +2024-07-09 13:39:26,946 INFO HandlerThread:1763125 [system_monitor.py:start():194] Starting system monitor +2024-07-09 13:39:26,946 INFO SystemMonitor:1763125 [system_monitor.py:_start():158] Starting system asset monitoring threads +2024-07-09 13:39:26,946 DEBUG SenderThread:1763125 [sender.py:send():379] send: stats +2024-07-09 13:39:26,946 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: pause +2024-07-09 13:39:26,947 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started cpu monitoring +2024-07-09 13:39:26,948 INFO HandlerThread:1763125 [handler.py:handle_request_pause():724] stopping system metrics thread +2024-07-09 13:39:26,948 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started disk monitoring +2024-07-09 13:39:26,948 INFO HandlerThread:1763125 [system_monitor.py:finish():203] Stopping system monitor +2024-07-09 13:39:26,948 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started gpu monitoring +2024-07-09 13:39:26,948 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():172] Starting system metrics aggregation loop +2024-07-09 13:39:26,949 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():179] Finished system metrics aggregation loop +2024-07-09 13:39:26,949 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():183] Publishing last batch of metrics +2024-07-09 13:39:26,951 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined cpu monitor +2024-07-09 13:39:26,952 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined disk monitor +2024-07-09 13:39:28,887 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined gpu monitor +2024-07-09 13:39:28,887 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: resume +2024-07-09 13:39:28,887 INFO HandlerThread:1763125 [handler.py:handle_request_resume():715] starting system metrics thread +2024-07-09 13:39:28,887 INFO HandlerThread:1763125 [system_monitor.py:start():194] Starting system monitor +2024-07-09 13:39:28,888 INFO SystemMonitor:1763125 [system_monitor.py:_start():158] Starting system asset monitoring threads +2024-07-09 13:39:28,888 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: pause +2024-07-09 13:39:28,888 INFO HandlerThread:1763125 [handler.py:handle_request_pause():724] stopping system metrics thread +2024-07-09 13:39:28,888 INFO HandlerThread:1763125 [system_monitor.py:finish():203] Stopping system monitor +2024-07-09 13:39:28,888 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started cpu monitoring +2024-07-09 13:39:28,888 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():172] Starting system metrics aggregation loop +2024-07-09 13:39:28,888 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():179] Finished system metrics aggregation loop +2024-07-09 13:39:28,888 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():183] Publishing last batch of metrics +2024-07-09 13:39:28,892 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined cpu monitor +2024-07-09 13:39:28,892 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: resume +2024-07-09 13:39:28,892 INFO HandlerThread:1763125 [handler.py:handle_request_resume():715] starting system metrics thread +2024-07-09 13:39:28,892 INFO HandlerThread:1763125 [system_monitor.py:start():194] Starting system monitor +2024-07-09 13:39:28,892 INFO SystemMonitor:1763125 [system_monitor.py:_start():158] Starting system asset monitoring threads +2024-07-09 13:39:28,892 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: internal_messages +2024-07-09 13:39:28,893 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started cpu monitoring +2024-07-09 13:39:28,893 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: pause +2024-07-09 13:39:28,893 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started disk monitoring +2024-07-09 13:39:28,893 INFO HandlerThread:1763125 [handler.py:handle_request_pause():724] stopping system metrics thread +2024-07-09 13:39:28,894 INFO HandlerThread:1763125 [system_monitor.py:finish():203] Stopping system monitor +2024-07-09 13:39:28,894 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started gpu monitoring +2024-07-09 13:39:28,894 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():172] Starting system metrics aggregation loop +2024-07-09 13:39:28,894 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():179] Finished system metrics aggregation loop +2024-07-09 13:39:28,894 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():183] Publishing last batch of metrics +2024-07-09 13:39:28,897 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined cpu monitor +2024-07-09 13:39:28,898 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined disk monitor +2024-07-09 13:39:30,786 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined gpu monitor +2024-07-09 13:39:30,786 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: resume +2024-07-09 13:39:30,786 INFO HandlerThread:1763125 [handler.py:handle_request_resume():715] starting system metrics thread +2024-07-09 13:39:30,786 INFO HandlerThread:1763125 [system_monitor.py:start():194] Starting system monitor +2024-07-09 13:39:30,786 INFO SystemMonitor:1763125 [system_monitor.py:_start():158] Starting system asset monitoring threads +2024-07-09 13:39:30,786 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: pause +2024-07-09 13:39:30,786 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started cpu monitoring +2024-07-09 13:39:30,787 INFO HandlerThread:1763125 [handler.py:handle_request_pause():724] stopping system metrics thread +2024-07-09 13:39:30,787 INFO HandlerThread:1763125 [system_monitor.py:finish():203] Stopping system monitor +2024-07-09 13:39:30,787 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started disk monitoring +2024-07-09 13:39:30,787 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():172] Starting system metrics aggregation loop +2024-07-09 13:39:30,787 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():179] Finished system metrics aggregation loop +2024-07-09 13:39:30,787 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():183] Publishing last batch of metrics +2024-07-09 13:39:30,790 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined cpu monitor +2024-07-09 13:39:30,790 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined disk monitor +2024-07-09 13:39:30,790 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: resume +2024-07-09 13:39:30,790 INFO HandlerThread:1763125 [handler.py:handle_request_resume():715] starting system metrics thread +2024-07-09 13:39:30,790 INFO HandlerThread:1763125 [system_monitor.py:start():194] Starting system monitor +2024-07-09 13:39:30,790 INFO SystemMonitor:1763125 [system_monitor.py:_start():158] Starting system asset monitoring threads +2024-07-09 13:39:30,790 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: pause +2024-07-09 13:39:30,791 DEBUG SenderThread:1763125 [sender.py:send():379] send: stats +2024-07-09 13:39:30,791 INFO HandlerThread:1763125 [handler.py:handle_request_pause():724] stopping system metrics thread +2024-07-09 13:39:30,791 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started cpu monitoring +2024-07-09 13:39:30,791 INFO HandlerThread:1763125 [system_monitor.py:finish():203] Stopping system monitor +2024-07-09 13:39:30,792 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started disk monitoring +2024-07-09 13:39:30,792 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():172] Starting system metrics aggregation loop +2024-07-09 13:39:30,792 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():179] Finished system metrics aggregation loop +2024-07-09 13:39:30,792 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():183] Publishing last batch of metrics +2024-07-09 13:39:30,795 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined cpu monitor +2024-07-09 13:39:30,796 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined disk monitor +2024-07-09 13:39:30,796 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:39:30,796 DEBUG SenderThread:1763125 [sender.py:send():379] send: stats +2024-07-09 13:39:30,797 DEBUG SenderThread:1763125 [sender.py:send():379] send: stats +2024-07-09 13:39:30,798 DEBUG SenderThread:1763125 [sender.py:send():379] send: stats +2024-07-09 13:39:30,801 DEBUG SenderThread:1763125 [sender.py:send():379] send: stats +2024-07-09 13:39:34,802 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:39:40,118 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: stop_status +2024-07-09 13:39:40,118 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: stop_status +2024-07-09 13:39:40,298 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:39:46,119 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:39:51,119 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:39:55,118 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: stop_status +2024-07-09 13:39:55,118 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: stop_status +2024-07-09 13:39:56,300 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:39:56,597 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/config.yaml +2024-07-09 13:40:02,120 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:08,120 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:10,118 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: stop_status +2024-07-09 13:40:10,118 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: stop_status +2024-07-09 13:40:13,850 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:19,120 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:24,120 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:25,118 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: stop_status +2024-07-09 13:40:25,118 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: stop_status +2024-07-09 13:40:29,299 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:35,120 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:40,118 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: stop_status +2024-07-09 13:40:40,119 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: stop_status +2024-07-09 13:40:40,297 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:46,120 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:51,120 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:40:55,118 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: stop_status +2024-07-09 13:40:55,119 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: stop_status +2024-07-09 13:40:56,297 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:41:02,121 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:41:07,121 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:41:07,491 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: resume +2024-07-09 13:41:07,491 INFO HandlerThread:1763125 [handler.py:handle_request_resume():715] starting system metrics thread +2024-07-09 13:41:07,491 INFO HandlerThread:1763125 [system_monitor.py:start():194] Starting system monitor +2024-07-09 13:41:07,491 INFO SystemMonitor:1763125 [system_monitor.py:_start():158] Starting system asset monitoring threads +2024-07-09 13:41:07,491 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started cpu monitoring +2024-07-09 13:41:07,492 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started disk monitoring +2024-07-09 13:41:07,492 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started gpu monitoring +2024-07-09 13:41:07,492 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started memory monitoring +2024-07-09 13:41:07,493 INFO SystemMonitor:1763125 [interfaces.py:start():188] Started network monitoring +2024-07-09 13:41:07,658 DEBUG SenderThread:1763125 [sender.py:send():379] send: telemetry +2024-07-09 13:41:07,689 DEBUG SenderThread:1763125 [sender.py:send():379] send: config +2024-07-09 13:41:07,721 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: server_info +2024-07-09 13:41:07,722 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: server_info +2024-07-09 13:41:07,871 DEBUG SenderThread:1763125 [sender.py:send():379] send: exit +2024-07-09 13:41:07,871 INFO SenderThread:1763125 [sender.py:send_exit():586] handling exit code: 0 +2024-07-09 13:41:07,871 INFO SenderThread:1763125 [sender.py:send_exit():588] handling runtime: 6 +2024-07-09 13:41:07,882 INFO SenderThread:1763125 [sender.py:_save_file():1454] saving file wandb-summary.json with policy end +2024-07-09 13:41:07,882 INFO SenderThread:1763125 [sender.py:send_exit():594] send defer +2024-07-09 13:41:07,883 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:07,883 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 0 +2024-07-09 13:41:07,883 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:07,883 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 0 +2024-07-09 13:41:07,883 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 1 +2024-07-09 13:41:07,883 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:07,883 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 1 +2024-07-09 13:41:07,883 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:07,883 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 1 +2024-07-09 13:41:07,883 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 2 +2024-07-09 13:41:07,883 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:07,883 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 2 +2024-07-09 13:41:07,883 INFO HandlerThread:1763125 [system_monitor.py:finish():203] Stopping system monitor +2024-07-09 13:41:07,884 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():172] Starting system metrics aggregation loop +2024-07-09 13:41:07,884 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined cpu monitor +2024-07-09 13:41:07,884 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():179] Finished system metrics aggregation loop +2024-07-09 13:41:07,885 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined disk monitor +2024-07-09 13:41:07,885 DEBUG SystemMonitor:1763125 [system_monitor.py:_start():183] Publishing last batch of metrics +2024-07-09 13:41:08,539 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json +2024-07-09 13:41:08,540 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log +2024-07-09 13:41:08,541 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/code/_session_history.ipynb +2024-07-09 13:41:08,541 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/code +2024-07-09 13:41:09,369 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined gpu monitor +2024-07-09 13:41:09,369 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined memory monitor +2024-07-09 13:41:09,369 INFO HandlerThread:1763125 [interfaces.py:finish():200] Joined network monitor +2024-07-09 13:41:09,369 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: poll_exit +2024-07-09 13:41:09,370 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:09,370 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 2 +2024-07-09 13:41:09,370 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 3 +2024-07-09 13:41:09,370 DEBUG SenderThread:1763125 [sender.py:send():379] send: stats +2024-07-09 13:41:09,370 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:09,371 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: poll_exit +2024-07-09 13:41:09,371 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 3 +2024-07-09 13:41:09,371 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:09,371 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 3 +2024-07-09 13:41:09,371 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 4 +2024-07-09 13:41:09,371 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:09,371 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 4 +2024-07-09 13:41:09,372 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:09,372 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 4 +2024-07-09 13:41:09,372 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 5 +2024-07-09 13:41:09,372 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:09,372 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 5 +2024-07-09 13:41:09,372 DEBUG SenderThread:1763125 [sender.py:send():379] send: summary +2024-07-09 13:41:09,399 INFO SenderThread:1763125 [sender.py:_save_file():1454] saving file wandb-summary.json with policy end +2024-07-09 13:41:09,400 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:09,400 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 5 +2024-07-09 13:41:09,400 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 6 +2024-07-09 13:41:09,400 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:09,400 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 6 +2024-07-09 13:41:09,400 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:09,400 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 6 +2024-07-09 13:41:09,402 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: status_report +2024-07-09 13:41:09,503 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 7 +2024-07-09 13:41:09,503 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:09,503 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 7 +2024-07-09 13:41:09,504 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:09,504 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 7 +2024-07-09 13:41:09,556 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/config.yaml +2024-07-09 13:41:09,556 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log +2024-07-09 13:41:09,556 INFO Thread-12 :1763125 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json +2024-07-09 13:41:09,871 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: poll_exit +2024-07-09 13:41:11,542 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 8 +2024-07-09 13:41:11,543 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: poll_exit +2024-07-09 13:41:11,543 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:11,543 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 8 +2024-07-09 13:41:11,543 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:11,543 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 8 +2024-07-09 13:41:11,543 INFO SenderThread:1763125 [job_builder.py:build():440] Attempting to build job artifact +2024-07-09 13:41:11,544 WARNING SenderThread:1763125 [job_builder.py:_log_if_verbose():274] Ensure read and write access to run files dir: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables +2024-07-09 13:41:11,544 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 9 +2024-07-09 13:41:11,544 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:11,544 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 9 +2024-07-09 13:41:11,544 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:11,544 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 9 +2024-07-09 13:41:11,544 INFO SenderThread:1763125 [dir_watcher.py:finish():358] shutting down directory watcher +2024-07-09 13:41:11,571 INFO SenderThread:1763125 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log +2024-07-09 13:41:11,575 INFO SenderThread:1763125 [dir_watcher.py:finish():388] scan: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files +2024-07-09 13:41:11,576 INFO SenderThread:1763125 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/requirements.txt requirements.txt +2024-07-09 13:41:11,576 INFO SenderThread:1763125 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log output.log +2024-07-09 13:41:11,576 INFO SenderThread:1763125 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/config.yaml config.yaml +2024-07-09 13:41:11,580 INFO SenderThread:1763125 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json wandb-summary.json +2024-07-09 13:41:11,586 INFO SenderThread:1763125 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/code/_session_history.ipynb code/_session_history.ipynb +2024-07-09 13:41:11,586 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 10 +2024-07-09 13:41:11,586 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:11,587 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 10 +2024-07-09 13:41:11,588 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:11,588 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 10 +2024-07-09 13:41:11,588 INFO SenderThread:1763125 [file_pusher.py:finish():169] shutting down file pusher +2024-07-09 13:41:11,874 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: poll_exit +2024-07-09 13:41:11,874 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: poll_exit +2024-07-09 13:41:11,979 INFO wandb-upload_0:1763125 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/requirements.txt +2024-07-09 13:41:12,141 INFO wandb-upload_1:1763125 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/output.log +2024-07-09 13:41:12,147 INFO wandb-upload_2:1763125 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/config.yaml +2024-07-09 13:41:12,182 INFO wandb-upload_3:1763125 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/wandb-summary.json +2024-07-09 13:41:12,210 INFO wandb-upload_4:1763125 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/files/code/_session_history.ipynb +2024-07-09 13:41:12,410 INFO Thread-11 (_thread_body):1763125 [sender.py:transition_state():614] send defer: 11 +2024-07-09 13:41:12,411 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:12,411 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 11 +2024-07-09 13:41:12,411 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:12,411 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 11 +2024-07-09 13:41:12,411 INFO SenderThread:1763125 [file_pusher.py:join():175] waiting for file pusher +2024-07-09 13:41:12,412 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 12 +2024-07-09 13:41:12,412 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:12,412 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 12 +2024-07-09 13:41:12,412 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:12,412 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 12 +2024-07-09 13:41:12,412 INFO SenderThread:1763125 [file_stream.py:finish():601] file stream finish called +2024-07-09 13:41:12,651 INFO SenderThread:1763125 [file_stream.py:finish():605] file stream finish is done +2024-07-09 13:41:12,651 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 13 +2024-07-09 13:41:12,652 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:12,652 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 13 +2024-07-09 13:41:12,652 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:12,652 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 13 +2024-07-09 13:41:12,652 INFO SenderThread:1763125 [sender.py:transition_state():614] send defer: 14 +2024-07-09 13:41:12,652 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: defer +2024-07-09 13:41:12,652 DEBUG SenderThread:1763125 [sender.py:send():379] send: final +2024-07-09 13:41:12,652 INFO HandlerThread:1763125 [handler.py:handle_request_defer():184] handle defer: 14 +2024-07-09 13:41:12,652 DEBUG SenderThread:1763125 [sender.py:send():379] send: footer +2024-07-09 13:41:12,652 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: defer +2024-07-09 13:41:12,652 INFO SenderThread:1763125 [sender.py:send_request_defer():610] handle sender defer: 14 +2024-07-09 13:41:12,653 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: poll_exit +2024-07-09 13:41:12,653 DEBUG SenderThread:1763125 [sender.py:send_request():406] send_request: poll_exit +2024-07-09 13:41:12,654 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: internal_messages +2024-07-09 13:41:12,654 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: get_summary +2024-07-09 13:41:12,655 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: sampled_history +2024-07-09 13:41:12,655 DEBUG HandlerThread:1763125 [handler.py:handle_request():158] handle_request: shutdown +2024-07-09 13:41:12,655 INFO HandlerThread:1763125 [handler.py:finish():882] shutting down handler +2024-07-09 13:41:13,653 INFO WriterThread:1763125 [datastore.py:close():296] close: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/run-semantic_cluster_0.1.wandb +2024-07-09 13:41:13,654 INFO SenderThread:1763125 [sender.py:finish():1608] shutting down sender +2024-07-09 13:41:13,654 INFO SenderThread:1763125 [file_pusher.py:finish():169] shutting down file pusher +2024-07-09 13:41:13,654 INFO SenderThread:1763125 [file_pusher.py:join():175] waiting for file pusher diff --git a/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug.log b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..b35df67cc795b975370e4474c63d1146903737bf --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug.log @@ -0,0 +1,75 @@ +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Current SDK version is 0.17.1 +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Configure stats pid to 1762217 +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Loading settings from /admin/home-ckadirt/.config/wandb/settings +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Loading settings from /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/settings +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Loading settings from environment variables: {} +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program': ''} +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Applying login settings: {'base_url': 'https://stability.wandb.io'} +2024-07-09 13:39:23,498 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Applying login settings: {} +2024-07-09 13:39:23,499 INFO MainThread:1762217 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug.log +2024-07-09 13:39:23,499 INFO MainThread:1762217 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/logs/debug-internal.log +2024-07-09 13:39:23,499 INFO MainThread:1762217 [wandb_init.py:_jupyter_setup():466] configuring jupyter hooks +2024-07-09 13:39:23,499 INFO MainThread:1762217 [wandb_init.py:init():560] calling init triggers +2024-07-09 13:39:23,500 INFO MainThread:1762217 [wandb_init.py:init():567] wandb.init called with sweep_config: {} +config: {'model_name': 'semantic_cluster_0.1', 'global_batch_size': 16, 'batch_size': 16, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 746793265, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 1e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'} +2024-07-09 13:39:23,500 INFO MainThread:1762217 [wandb_init.py:init():610] starting backend +2024-07-09 13:39:23,500 INFO MainThread:1762217 [wandb_init.py:init():614] setting up manager +2024-07-09 13:39:23,510 INFO MainThread:1762217 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn +2024-07-09 13:39:23,517 INFO MainThread:1762217 [wandb_init.py:init():622] backend started and connected +2024-07-09 13:39:23,539 INFO MainThread:1762217 [wandb_run.py:_label_probe_notebook():1334] probe notebook +2024-07-09 13:39:23,541 INFO MainThread:1762217 [wandb_run.py:_label_probe_notebook():1344] Unable to probe notebook: 'NoneType' object has no attribute 'get' +2024-07-09 13:39:23,541 INFO MainThread:1762217 [wandb_init.py:init():711] updated telemetry +2024-07-09 13:39:23,659 INFO MainThread:1762217 [wandb_init.py:init():744] communicating run to backend with 90.0 second timeout +2024-07-09 13:39:24,183 INFO MainThread:1762217 [wandb_init.py:init():787] run resumed +2024-07-09 13:39:24,219 INFO MainThread:1762217 [wandb_run.py:_on_init():2402] communicating current version +2024-07-09 13:39:24,273 INFO MainThread:1762217 [wandb_run.py:_on_init():2411] got version response upgrade_message: "wandb version 0.17.4 is available! To upgrade, please run:\n $ pip install wandb --upgrade" + +2024-07-09 13:39:24,273 INFO MainThread:1762217 [wandb_init.py:init():795] starting run threads in backend +2024-07-09 13:39:25,118 INFO MainThread:1762217 [wandb_run.py:_console_start():2380] atexit reg +2024-07-09 13:39:25,118 INFO MainThread:1762217 [wandb_run.py:_redirect():2235] redirect: wrap_raw +2024-07-09 13:39:25,118 INFO MainThread:1762217 [wandb_run.py:_redirect():2300] Wrapping output streams. +2024-07-09 13:39:25,119 INFO MainThread:1762217 [wandb_run.py:_redirect():2325] Redirects installed. +2024-07-09 13:39:25,126 INFO MainThread:1762217 [wandb_init.py:init():838] run started, returning control to user process +2024-07-09 13:39:25,132 INFO MainThread:1762217 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 13:39:25,132 INFO MainThread:1762217 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 13:39:25,246 INFO MainThread:1762217 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 13:39:25,247 INFO MainThread:1762217 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 13:39:25,248 INFO MainThread:1762217 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 13:39:25,320 INFO MainThread:1762217 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 13:39:25,321 INFO MainThread:1762217 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 13:39:25,322 INFO MainThread:1762217 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 13:39:25,389 INFO MainThread:1762217 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 13:39:26,503 INFO MainThread:1762217 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 13:39:26,504 INFO MainThread:1762217 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 13:39:26,631 INFO MainThread:1762217 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 13:39:26,632 INFO MainThread:1762217 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 13:39:26,633 INFO MainThread:1762217 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 13:39:26,705 INFO MainThread:1762217 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 13:39:27,023 INFO MainThread:1762217 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 13:39:27,023 INFO MainThread:1762217 [wandb_init.py:_pause_backend():431] pausing backend +2024-07-09 13:41:07,490 INFO MainThread:1762217 [wandb_init.py:_resume_backend():436] resuming backend +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Current SDK version is 0.17.1 +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Configure stats pid to 1762217 +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Loading settings from /admin/home-ckadirt/.config/wandb/settings +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Loading settings from /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/settings +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Loading settings from environment variables: {} +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program': ''} +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Applying login settings: {'base_url': 'https://stability.wandb.io'} +2024-07-09 13:41:07,654 INFO MainThread:1762217 [wandb_setup.py:_flush():76] Applying login settings: {} +2024-07-09 13:41:07,655 INFO MainThread:1762217 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_134107-semantic_cluster_0.1/logs/debug.log +2024-07-09 13:41:07,656 INFO MainThread:1762217 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_134107-semantic_cluster_0.1/logs/debug-internal.log +2024-07-09 13:41:07,656 INFO MainThread:1762217 [wandb_init.py:init():560] calling init triggers +2024-07-09 13:41:07,656 INFO MainThread:1762217 [wandb_init.py:init():567] wandb.init called with sweep_config: {} +config: {'model_name': 'semantic_cluster_0.1', 'global_batch_size': 16, 'batch_size': 16, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 746793265, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 1e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'} +2024-07-09 13:41:07,656 INFO MainThread:1762217 [wandb_init.py:init():585] re-initializing run, found existing run on stack: semantic_cluster_0.1 +2024-07-09 13:41:07,658 INFO MainThread:1762217 [wandb_run.py:_finish():2109] finishing run ckadirt/mindeye_semantic_cluster_0.1/semantic_cluster_0.1 +2024-07-09 13:41:07,688 INFO MainThread:1762217 [jupyter.py:save_history():473] saving 27 cells to _session_history.ipynb +2024-07-09 13:41:07,688 INFO MainThread:1762217 [wandb_run.py:_config_callback():1382] config_cb ('_wandb', 'session_history') code/_session_history.ipynb None +2024-07-09 13:41:07,719 INFO MainThread:1762217 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/Untitled1.ipynb +2024-07-09 13:41:07,719 INFO MainThread:1762217 [wandb_init.py:_jupyter_teardown():448] cleaning up jupyter logic +2024-07-09 13:41:07,720 INFO MainThread:1762217 [wandb_run.py:_atexit_cleanup():2349] got exitcode: 0 +2024-07-09 13:41:07,721 INFO MainThread:1762217 [wandb_run.py:_restore():2332] restore +2024-07-09 13:41:07,721 INFO MainThread:1762217 [wandb_run.py:_restore():2338] restore done +2024-07-09 13:41:13,657 INFO MainThread:1762217 [wandb_run.py:_footer_history_summary_info():4008] rendering history +2024-07-09 13:41:13,657 INFO MainThread:1762217 [wandb_run.py:_footer_history_summary_info():4040] rendering summary +2024-07-09 13:41:13,676 INFO MainThread:1762217 [wandb_run.py:_footer_sync_info():3967] logging synced files diff --git a/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/tmp/code/_session_history.ipynb b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/tmp/code/_session_history.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..96db24afeea6d5848084bdab06c6b5a38f80175d --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1/tmp/code/_session_history.ipynb @@ -0,0 +1,1036 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "770e88fe", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import sys\n", + "import json\n", + "import argparse\n", + "import numpy as np\n", + "import math\n", + "from einops import rearrange\n", + "import time\n", + "import random\n", + "import string\n", + "import h5py\n", + "from tqdm import tqdm\n", + "import webdataset as wds\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import torch\n", + "import torch.nn as nn\n", + "from torchvision import transforms\n", + "from accelerate import Accelerator\n", + "import torch.nn.functional as F\n", + "\n", + "# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main\n", + "sys.path.append('generative_models/')\n", + "import sgm\n", + "from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder\n", + "\n", + "# tf32 data type is faster than standard float32\n", + "torch.backends.cuda.matmul.allow_tf32 = True\n", + "\n", + "# custom functions #\n", + "import utils" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "38d13e72", + "metadata": {}, + "outputs": [], + "source": [ + "def classPrecision(logits, y_true, top=1):\n", + " \"\"\"\n", + " Calculate the precision of the top-n predictions.\n", + " \n", + " Parameters:\n", + " logits (torch.Tensor): The output logits from the model (shape: [batch_size, num_classes]).\n", + " y_true (torch.Tensor): The ground truth labels (shape: [batch_size]).\n", + " top (int): The number of top predictions to consider.\n", + " \n", + " Returns:\n", + " float: The precision percentage of the top-n predictions.\n", + " \"\"\"\n", + " # Apply softmax to get probabilities\n", + " probs = F.softmax(logits, dim=1).detach().cpu()\n", + " \n", + " # Get the top-n predictions\n", + " top_n_preds = torch.topk(probs, top, dim=1).indices.detach().cpu()\n", + "\n", + " # Move y_true to CPU and detach\n", + " y_true = y_true.detach().cpu()\n", + "\n", + " # Check if y_true is in top-n predictions\n", + " correct = top_n_preds.eq(y_true.view(-1, 1).expand_as(top_n_preds))\n", + "\n", + " # Calculate precision\n", + " precision = correct.sum().item() / y_true.size(0)\n", + " \n", + " return precision * 100\n", + "\n", + "# Example usage:\n", + "logits = torch.randn(8, 41) # Example logits tensor\n", + "y_true = torch.randint(0, 41, (8,)) # Example ground truth labels\n", + "\n", + "top_n_precision = classPrecision(logits, y_true, top=1)\n", + "print(f\"Top-1 Precision: {top_n_precision:.2f}%\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f8410934", + "metadata": {}, + "outputs": [], + "source": [ + "### Multi-GPU config ###\n", + "local_rank = os.getenv('RANK')\n", + "if local_rank is None: \n", + " local_rank = 0\n", + "else:\n", + " local_rank = int(local_rank)\n", + "print(\"LOCAL RANK \", local_rank) \n", + "\n", + "data_type = torch.float16 # change depending on your mixed_precision\n", + "num_devices = torch.cuda.device_count()\n", + "if num_devices==0: num_devices = 1\n", + "\n", + "# First use \"accelerate config\" in terminal and setup using deepspeed stage 2 with CPU offloading!\n", + "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n", + "if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job\n", + " global_batch_size = batch_size = 16\n", + "else:\n", + " global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n", + " batch_size = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a02d706c", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"PID of this process =\",os.getpid())\n", + "device = accelerator.device\n", + "print(\"device:\",device)\n", + "world_size = accelerator.state.num_processes\n", + "distributed = not accelerator.state.distributed_type == 'NO'\n", + "num_devices = torch.cuda.device_count()\n", + "if num_devices==0 or not distributed: num_devices = 1\n", + "num_workers = num_devices\n", + "print(accelerator.state)\n", + "\n", + "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\n", + "print = accelerator.print # only print if local_rank=0" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ab031a50", + "metadata": {}, + "outputs": [], + "source": [ + "# if running this interactively, can specify jupyter_args here for argparser to use\n", + "if utils.is_interactive():\n", + " model_name = \"semantic_cluster_0.1\"\n", + " print(\"model_name:\", model_name)\n", + " \n", + " # global_batch_size and batch_size should already be defined in the 2nd cell block\n", + " jupyter_args = f\"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \\\n", + " --cache_dir=/weka/proj-medarc/shared/cache \\\n", + " --model_name={model_name} \\\n", + " --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \\\n", + " --hidden_dim=1024 --clip_scale=1. \\\n", + " --no-blurry_recon --blur_scale=.5 \\\n", + " --use_prior --prior_scale=30 \\\n", + " --n_blocks=4 --max_lr=1e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \\\n", + " --ckpt_interval=999 --no-ckpt_saving --wandb_log\"\n", + " # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow\n", + "\n", + " print(jupyter_args)\n", + " jupyter_args = jupyter_args.split()\n", + " \n", + " from IPython.display import clear_output # function to clear print outputs in cell\n", + " %load_ext autoreload \n", + " # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions\n", + " %autoreload 2 " + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "afbef21c", + "metadata": {}, + "outputs": [], + "source": [ + "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n", + "parser.add_argument(\n", + " \"--model_name\", type=str, default=\"testing2\",\n", + " help=\"name of model, used for ckpt saving and wandb logging (if enabled)\",\n", + ")\n", + "parser.add_argument(\n", + " \"--data_path\", type=str, default=os.getcwd(),\n", + " help=\"Path to where NSD data is stored / where to download it to\",\n", + ")\n", + "parser.add_argument(\n", + " \"--cache_dir\", type=str, default=os.getcwd(),\n", + " help=\"Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.\",\n", + ")\n", + "parser.add_argument(\n", + " \"--subj\",type=int, default=1, choices=[1,2,3,4,5,6,7,8],\n", + " help=\"Validate on which subject?\",\n", + ")\n", + "parser.add_argument(\n", + " \"--multisubject_ckpt\", type=str, default=None,\n", + " help=\"Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.\",\n", + ")\n", + "parser.add_argument(\n", + " \"--num_sessions\", type=int, default=1,\n", + " help=\"Number of training sessions to include\",\n", + ")\n", + "parser.add_argument(\n", + " \"--use_prior\",action=argparse.BooleanOptionalAction,default=True,\n", + " help=\"whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)\",\n", + ")\n", + "parser.add_argument(\n", + " \"--batch_size\", type=int, default=16,\n", + " help=\"Batch size can be increased by 10x if only training retreival submodule and not diffusion prior\",\n", + ")\n", + "parser.add_argument(\n", + " \"--wandb_log\",action=argparse.BooleanOptionalAction,default=False,\n", + " help=\"whether to log to wandb\",\n", + ")\n", + "parser.add_argument(\n", + " \"--wandb_project\",type=str,default=\"stability\",\n", + " help=\"wandb project name\",\n", + ")\n", + "parser.add_argument(\n", + " \"--mixup_pct\",type=float,default=.33,\n", + " help=\"proportion of way through training when to switch from BiMixCo to SoftCLIP\",\n", + ")\n", + "parser.add_argument(\n", + " \"--blurry_recon\",action=argparse.BooleanOptionalAction,default=True,\n", + " help=\"whether to output blurry reconstructions\",\n", + ")\n", + "parser.add_argument(\n", + " \"--blur_scale\",type=float,default=.5,\n", + " help=\"multiply loss from blurry recons by this number\",\n", + ")\n", + "parser.add_argument(\n", + " \"--clip_scale\",type=float,default=1.,\n", + " help=\"multiply contrastive loss by this number\",\n", + ")\n", + "parser.add_argument(\n", + " \"--prior_scale\",type=float,default=30,\n", + " help=\"multiply diffusion prior loss by this\",\n", + ")\n", + "parser.add_argument(\n", + " \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=False,\n", + " help=\"whether to use image augmentation\",\n", + ")\n", + "parser.add_argument(\n", + " \"--num_epochs\",type=int,default=150,\n", + " help=\"number of epochs of training\",\n", + ")\n", + "parser.add_argument(\n", + " \"--multi_subject\",action=argparse.BooleanOptionalAction,default=False,\n", + ")\n", + "parser.add_argument(\n", + " \"--new_test\",action=argparse.BooleanOptionalAction,default=True,\n", + ")\n", + "parser.add_argument(\n", + " \"--n_blocks\",type=int,default=4,\n", + ")\n", + "parser.add_argument(\n", + " \"--hidden_dim\",type=int,default=1024,\n", + ")\n", + "parser.add_argument(\n", + " \"--lr_scheduler_type\",type=str,default='cycle',choices=['cycle','linear'],\n", + ")\n", + "parser.add_argument(\n", + " \"--ckpt_saving\",action=argparse.BooleanOptionalAction,default=True,\n", + ")\n", + "parser.add_argument(\n", + " \"--ckpt_interval\",type=int,default=5,\n", + " help=\"save backup ckpt and reconstruct every x epochs\",\n", + ")\n", + "parser.add_argument(\n", + " \"--seed\",type=int,default=42,\n", + ")\n", + "parser.add_argument(\n", + " \"--max_lr\",type=float,default=3e-5,\n", + ")\n", + "\n", + "if utils.is_interactive():\n", + " args = parser.parse_args(jupyter_args)\n", + "else:\n", + " args = parser.parse_args()\n", + "\n", + "# create global variables without the args prefix\n", + "for attribute_name in vars(args).keys():\n", + " globals()[attribute_name] = getattr(args, attribute_name)\n", + " \n", + "# seed all random functions\n", + "utils.seed_everything(seed)\n", + "\n", + "outdir = os.path.abspath(f'../train_logs/{model_name}')\n", + "if not os.path.exists(outdir) and ckpt_saving:\n", + " os.makedirs(outdir,exist_ok=True)\n", + " \n", + "if use_image_aug or blurry_recon:\n", + " import kornia\n", + " from kornia.augmentation.container import AugmentationSequential\n", + "if use_image_aug:\n", + " img_augment = AugmentationSequential(\n", + " kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),\n", + " same_on_batch=False,\n", + " data_keys=[\"input\"],\n", + " )\n", + " \n", + "if multi_subject:\n", + " subj_list = np.arange(1,9)\n", + " subj_list = subj_list[subj_list != subj]\n", + "else:\n", + " subj_list = [subj]\n", + "\n", + "print(\"subj_list\", subj_list, \"num_sessions\", num_sessions)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "636f61c6", + "metadata": {}, + "outputs": [], + "source": [ + "def my_split_by_node(urls): return urls\n", + "num_voxels_list = []\n", + "\n", + "if multi_subject:\n", + " nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])\n", + " num_samples_per_epoch = (750*40) // num_devices \n", + "else:\n", + " num_samples_per_epoch = (750*num_sessions) // num_devices \n", + "\n", + "print(\"dividing batch size by subj_list, which will then be concatenated across subj during training...\") \n", + "batch_size = batch_size // len(subj_list)\n", + "\n", + "num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))\n", + "\n", + "print(\"batch_size =\", batch_size, \"num_iterations_per_epoch =\",num_iterations_per_epoch, \"num_samples_per_epoch =\",num_samples_per_epoch)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "55dcf898", + "metadata": {}, + "outputs": [], + "source": [ + "train_data = {}\n", + "train_dl = {}\n", + "num_voxels = {}\n", + "voxels = {}\n", + "for s in subj_list:\n", + " print(f\"Training with {num_sessions} sessions\")\n", + " if multi_subject:\n", + " train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{nsessions_allsubj[s-1]-1}\" + \"}.tar\"\n", + " else:\n", + " train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n", + " print(train_url)\n", + " \n", + " train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,nodesplitter=my_split_by_node)\\\n", + " .shuffle(750, initial=1500, rng=random.Random(42))\\\n", + " .decode(\"torch\")\\\n", + " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", + " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", + " train_dl[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n", + "\n", + " f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n", + " betas = f['betas'][:]\n", + " betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n", + " num_voxels_list.append(betas[0].shape[-1])\n", + " num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n", + " voxels[f'subj0{s}'] = betas\n", + " print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n", + "\n", + "print(\"Loaded all subj train dls and betas!\\n\")\n", + "\n", + "# Validate only on one subject\n", + "if multi_subject: \n", + " subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list\n", + "if not new_test: # using old test set from before full dataset released (used in original MindEye paper)\n", + " if subj==3:\n", + " num_test=2113\n", + " elif subj==4:\n", + " num_test=1985\n", + " elif subj==6:\n", + " num_test=2113\n", + " elif subj==8:\n", + " num_test=1985\n", + " else:\n", + " num_test=2770\n", + " test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n", + "elif new_test: # using larger test set from after full dataset released\n", + " if subj==3:\n", + " num_test=2371\n", + " elif subj==4:\n", + " num_test=2188\n", + " elif subj==6:\n", + " num_test=2371\n", + " elif subj==8:\n", + " num_test=2188\n", + " else:\n", + " num_test=3000\n", + " test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n", + "print(test_url)\n", + "test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\\\n", + " .shuffle(750, initial=1500, rng=random.Random(42))\\\n", + " .decode(\"torch\")\\\n", + " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", + " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", + "test_dl = torch.utils.data.DataLoader(test_data, batch_size=num_test, shuffle=False, drop_last=True, pin_memory=True)\n", + "print(f\"Loaded test dl for subj{subj}!\\n\")" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "3ac1391b", + "metadata": {}, + "outputs": [], + "source": [ + "# Load 73k NSD images\n", + "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n", + "images = f['images']\n", + "print(\"Loaded all 73k possible NSD images to cpu!\", images.shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "4f4acf44", + "metadata": {}, + "outputs": [], + "source": [ + "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n", + " arch=\"ViT-bigG-14\",\n", + " version=\"laion2b_s39b_b160k\",\n", + " output_tokens=True,\n", + " only_tokens=True,\n", + ")\n", + "clip_img_embedder.to(device)\n", + "\n", + "clip_seq_dim = 256\n", + "clip_emb_dim = 1664" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "9bc03c5e", + "metadata": {}, + "outputs": [], + "source": [ + "if blurry_recon:\n", + " from diffusers import AutoencoderKL \n", + " autoenc = AutoencoderKL(\n", + " down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],\n", + " up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],\n", + " block_out_channels=[128, 256, 512, 512],\n", + " layers_per_block=2,\n", + " sample_size=256,\n", + " )\n", + " ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')\n", + " autoenc.load_state_dict(ckpt)\n", + " \n", + " autoenc.eval()\n", + " autoenc.requires_grad_(False)\n", + " autoenc.to(device)\n", + " utils.count_params(autoenc)\n", + " \n", + " from autoencoder.convnext import ConvnextXL\n", + " cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')\n", + " cnx.requires_grad_(False)\n", + " cnx.eval()\n", + " cnx.to(device)\n", + " \n", + " mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)\n", + " std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)\n", + " \n", + " blur_augs = AugmentationSequential(\n", + " kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),\n", + " kornia.augmentation.RandomGrayscale(p=0.1),\n", + " kornia.augmentation.RandomSolarize(p=0.1),\n", + " kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),\n", + " data_keys=[\"input\"],\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "fbf1e6fb", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "MindEyeModule()" + ] + } + ], + "source": [ + "class MindEyeModule(nn.Module):\n", + " def __init__(self):\n", + " super(MindEyeModule, self).__init__()\n", + " def forward(self, x):\n", + " return x\n", + " \n", + "model = MindEyeModule()\n", + "model" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "f291a3f5", + "metadata": {}, + "outputs": [], + "source": [ + "class RidgeRegression(torch.nn.Module):\n", + " # make sure to add weight_decay when initializing optimizer to enable regularization\n", + " def __init__(self, input_sizes, out_features): \n", + " super(RidgeRegression, self).__init__()\n", + " self.out_features = out_features\n", + " self.linears = torch.nn.ModuleList([\n", + " torch.nn.Linear(input_size, out_features) for input_size in input_sizes\n", + " ])\n", + " def forward(self, x, subj_idx):\n", + " out = self.linears[subj_idx](x[:,0]).unsqueeze(1)\n", + " return out\n", + " \n", + "class IndividRidgeRegression(torch.nn.Module):\n", + " def __init__(self, input_size, out_features):\n", + " super(IndividRidgeRegression, self).__init__()\n", + " self.out_features = out_features\n", + " self.linear = torch.nn.Linear(input_size, out_features)\n", + " def forward(self, x):\n", + " out = self.linear(x)\n", + " return out\n", + " \n", + "model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)\n", + "utils.count_params(model.ridge)\n", + "utils.count_params(model)\n", + "\n", + "# test on subject 1 with fake data\n", + "b = torch.randn((2,1,num_voxels_list[0]))\n", + "print(b.shape, model.ridge(b,0).shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "1eec72cb", + "metadata": {}, + "outputs": [], + "source": [ + "from models import BrainNetwork\n", + "model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,\n", + " clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, \n", + " blurry_recon=blurry_recon, clip_scale=clip_scale)\n", + "utils.count_params(model.backbone)\n", + "utils.count_params(model)\n", + "\n", + "# test that the model works on some fake data\n", + "b = torch.randn((2,1,hidden_dim))\n", + "print(\"b.shape\",b.shape)\n", + "\n", + "backbone_, clip_, blur_ = model.backbone(b)\n", + "print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "f72a54e8", + "metadata": {}, + "outputs": [], + "source": [ + "if use_prior:\n", + " from models import *\n", + "\n", + " # setup diffusion prior network\n", + " out_dim = clip_emb_dim\n", + " depth = 6\n", + " dim_head = 52\n", + " heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim\n", + " timesteps = 100\n", + "\n", + " prior_network = PriorNetwork(\n", + " dim=out_dim,\n", + " depth=depth,\n", + " dim_head=dim_head,\n", + " heads=heads,\n", + " causal=False,\n", + " num_tokens = clip_seq_dim,\n", + " learned_query_mode=\"pos_emb\"\n", + " )\n", + "\n", + " model.diffusion_prior = BrainDiffusionPrior(\n", + " net=prior_network,\n", + " image_embed_dim=out_dim,\n", + " condition_on_text_encodings=False,\n", + " timesteps=timesteps,\n", + " cond_drop_prob=0.2,\n", + " image_embed_scale=None,\n", + " )\n", + " \n", + " utils.count_params(model.diffusion_prior)\n", + " utils.count_params(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "57803a48", + "metadata": {}, + "outputs": [], + "source": [ + "path_semantic_names = \"/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy\"\n", + "path_semantic_cluster = \"/weka/proj-fmri/ckadirt/MindEyeV2/src/COCO_73k_semantic_cluster.npy\"\n", + "semantic_cluster_names = np.load(path_semantic_names)\n", + "semantic_cluster = np.load(path_semantic_cluster)\n", + "possible_semantic_clusters = np.unique(semantic_cluster)\n", + "\n", + "# one-hot encode semantic clusters\n", + "# move possible_semantic_clusters to numbers and create a dictionary\n", + "semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)}\n", + "semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters)))\n", + "for i, cluster in enumerate(semantic_cluster):\n", + " semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1\n", + "\n", + "\n", + "print(\"semantic_cluster_onehot.shape\", semantic_cluster_onehot.shape)\n", + "\n", + "num_seman_clusters = len(np.unique(semantic_cluster))\n", + "print(\"num_seman_clusters\", num_seman_clusters)" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "4497c76c", + "metadata": {}, + "outputs": [], + "source": [ + "# plot some images next to their semantic cluster\n", + "fig, ax = plt.subplots(1, 5, figsize=(20, 4))\n", + "for i in range(5):\n", + " # covert numpy array images to float32 \n", + " image_index = torch.randint(0, len(images), (1,)).item()\n", + " print(image_index)\n", + " ax[i].imshow(images[image_index].transpose(1,2,0).astype(np.float32))\n", + " ax[i].set_title(semantic_cluster[image_index])\n", + " ax[i].axis(\"off\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "754f285d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "746793265" + ] + } + ], + "source": [ + "model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters)\n", + "utils.count_params(model.RRClassifier)\n", + "utils.count_params(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "7be4a60e", + "metadata": {}, + "outputs": [], + "source": [ + "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n", + "\n", + "opt_grouped_parameters = [\n", + " {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},\n", + " {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n", + " {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n", + " {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1e-2},\n", + "]\n", + "# if use_prior:\n", + "# opt_grouped_parameters.extend([\n", + "# {'params': [p for n, p in model.diffusion_prior.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n", + "# {'params': [p for n, p in model.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n", + "# ])\n", + "# opt_grouped_parameters.extend([\n", + "# \n", + "# ])\n", + "\n", + "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)\n", + "\n", + "if lr_scheduler_type == 'linear':\n", + " lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n", + " optimizer,\n", + " total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),\n", + " last_epoch=-1\n", + " )\n", + "elif lr_scheduler_type == 'cycle':\n", + " total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))\n", + " print(\"total_steps\", total_steps)\n", + " lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n", + " optimizer, \n", + " max_lr=max_lr,\n", + " total_steps=total_steps,\n", + " final_div_factor=1000,\n", + " last_epoch=-1, pct_start=2/num_epochs\n", + " )\n", + " \n", + "def save_ckpt(tag):\n", + " ckpt_path = outdir+f'/{tag}.pth'\n", + " if accelerator.is_main_process:\n", + " unwrapped_model = accelerator.unwrap_model(model)\n", + " torch.save({\n", + " 'epoch': epoch,\n", + " 'model_state_dict': unwrapped_model.state_dict(),\n", + " 'optimizer_state_dict': optimizer.state_dict(),\n", + " 'lr_scheduler': lr_scheduler.state_dict(),\n", + " 'train_losses': losses,\n", + " 'test_losses': test_losses,\n", + " 'lrs': lrs,\n", + " }, ckpt_path)\n", + " print(f\"\\n---saved {outdir}/{tag} ckpt!---\\n\")\n", + "\n", + "def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): \n", + " print(f\"\\n---loading {outdir}/{tag}.pth ckpt---\\n\")\n", + " checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n", + " state_dict = checkpoint['model_state_dict']\n", + " if multisubj_loading: # remove incompatible ridge layer that will otherwise error\n", + " state_dict.pop('ridge.linears.0.weight',None)\n", + " model.load_state_dict(state_dict, strict=strict)\n", + " if load_epoch:\n", + " globals()[\"epoch\"] = checkpoint['epoch']\n", + " print(\"Epoch\",epoch)\n", + " if load_optimizer:\n", + " optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n", + " if load_lr:\n", + " lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n", + " del checkpoint\n", + "\n", + "print(\"\\nDone with model preparations!\")\n", + "num_params = utils.count_params(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "6fedf0ad", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1e-05" + ] + } + ], + "source": [ + "max_lr" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "7d78e7f3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "wandb version 0.17.4 is available! To upgrade, please run:\n", + " $ pip install wandb --upgrade" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Tracking run with wandb version 0.17.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Run data is saved locally in /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_133923-semantic_cluster_0.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Resuming run semantic_cluster_0.1 to Weights & Biases (docs)
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + " View project at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster_0.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + " View run at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster_0.1/runs/semantic_cluster_0.1" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "if local_rank==0 and wandb_log: # only use main process for wandb logging\n", + " import wandb\n", + " wandb_project = 'mindeye_semantic_cluster_0.1'\n", + " print(f\"wandb {wandb_project} run {model_name}\")\n", + " # need to configure wandb beforehand in terminal with \"wandb init\"!\n", + " wandb_config = {\n", + " \"model_name\": model_name,\n", + " \"global_batch_size\": global_batch_size,\n", + " \"batch_size\": batch_size,\n", + " \"num_epochs\": num_epochs,\n", + " \"num_sessions\": num_sessions,\n", + " \"num_params\": num_params,\n", + " \"clip_scale\": clip_scale,\n", + " \"prior_scale\": prior_scale,\n", + " \"blur_scale\": blur_scale,\n", + " \"use_image_aug\": use_image_aug,\n", + " \"max_lr\": max_lr,\n", + " \"mixup_pct\": mixup_pct,\n", + " \"num_samples_per_epoch\": num_samples_per_epoch,\n", + " \"num_test\": num_test,\n", + " \"ckpt_interval\": ckpt_interval,\n", + " \"ckpt_saving\": ckpt_saving,\n", + " \"seed\": seed,\n", + " \"distributed\": distributed,\n", + " \"num_devices\": num_devices,\n", + " \"world_size\": world_size,\n", + " \"train_url\": train_url,\n", + " \"test_url\": test_url,\n", + " }\n", + " print(\"wandb_config:\\n\",wandb_config)\n", + " print(\"wandb_id:\",model_name)\n", + " wandb.login(host='https://stability.wandb.io')\n", + " wandb.init(\n", + " id=model_name,\n", + " project=wandb_project,\n", + " name=model_name,\n", + " config=wandb_config,\n", + " resume=\"allow\",\n", + " )\n", + "else:\n", + " wandb_log = False" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "69011870", + "metadata": {}, + "outputs": [], + "source": [ + "epoch = 0\n", + "losses, test_losses, lrs = [], [], []\n", + "best_test_loss = 1e9\n", + "torch.cuda.empty_cache()" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "cbc20574", + "metadata": {}, + "outputs": [], + "source": [ + "# load multisubject stage1 ckpt if set\n", + "if multisubject_ckpt is not None:\n", + " load_ckpt(\"last\",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "d89584ca", + "metadata": {}, + "outputs": [], + "source": [ + "train_dls = [train_dl[f'subj0{s}'] for s in subj_list]\n", + "\n", + "model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot)\n", + "# leaving out test_dl since we will only have local_rank 0 device do evals" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "47b1e63b", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_semantic_clusters(images, indexes, semantic_cluster, semantic_cluster_dict):\n", + " # check images and indexes are the same length\n", + " assert len(images) == len(indexes)\n", + " \n", + " # if images are tensors, convert them to numpy arrays and move them to cpu\n", + " if isinstance(images, torch.Tensor):\n", + " images = images.cpu().numpy()\n", + "\n", + " fig, ax = plt.subplots(1, len(images), figsize=(20, 4))\n", + " for i, index in enumerate(indexes):\n", + " # covert numpy array images to float32 \n", + " ax[i].imshow(images[i].transpose(1,2,0).astype(np.float32))\n", + " # search the key in the dictionary based on the index as value\n", + " name = {i for i in semantic_cluster_dict if semantic_cluster_dict[i] == index}\n", + " ax[i].set_title(name)\n", + " ax[i].axis(\"off\")\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "fd22d30b", + "metadata": {}, + "outputs": [], + "source": [ + "RTTT" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "d351dffe", + "metadata": {}, + "outputs": [], + "source": [ + "if local_rank==0 and wandb_log: # only use main process for wandb logging\n", + " import wandb\n", + " wandb_project = 'mindeye_semantic_cluster_0.2'\n", + " print(f\"wandb {wandb_project} run {model_name}\")\n", + " # need to configure wandb beforehand in terminal with \"wandb init\"!\n", + " wandb_config = {\n", + " \"model_name\": model_name,\n", + " \"global_batch_size\": global_batch_size,\n", + " \"batch_size\": batch_size,\n", + " \"num_epochs\": num_epochs,\n", + " \"num_sessions\": num_sessions,\n", + " \"num_params\": num_params,\n", + " \"clip_scale\": clip_scale,\n", + " \"prior_scale\": prior_scale,\n", + " \"blur_scale\": blur_scale,\n", + " \"use_image_aug\": use_image_aug,\n", + " \"max_lr\": max_lr,\n", + " \"mixup_pct\": mixup_pct,\n", + " \"num_samples_per_epoch\": num_samples_per_epoch,\n", + " \"num_test\": num_test,\n", + " \"ckpt_interval\": ckpt_interval,\n", + " \"ckpt_saving\": ckpt_saving,\n", + " \"seed\": seed,\n", + " \"distributed\": distributed,\n", + " \"num_devices\": num_devices,\n", + " \"world_size\": world_size,\n", + " \"train_url\": train_url,\n", + " \"test_url\": test_url,\n", + " }\n", + " print(\"wandb_config:\\n\",wandb_config)\n", + " print(\"wandb_id:\",model_name)\n", + " wandb.login(host='https://stability.wandb.io')\n", + " wandb.init(\n", + " id=model_name,\n", + " project=wandb_project,\n", + " name=model_name,\n", + " config=wandb_config,\n", + " resume=\"allow\",\n", + " )\n", + "else:\n", + " wandb_log = False" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "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.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/code/src/Untitled1.py b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/code/src/Untitled1.py new file mode 100644 index 0000000000000000000000000000000000000000..961b602e024548bb13741bafa7a86e85592bf1b8 --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/code/src/Untitled1.py @@ -0,0 +1,1161 @@ +#!/usr/bin/env python +# coding: utf-8 + +# In[1]: + + +import os +import sys +import json +import argparse +import numpy as np +import math +from einops import rearrange +import time +import random +import string +import h5py +from tqdm import tqdm +import webdataset as wds + +import matplotlib.pyplot as plt +import torch +import torch.nn as nn +from torchvision import transforms +from accelerate import Accelerator +import torch.nn.functional as F + +# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main +sys.path.append('generative_models/') +import sgm +from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder + +# tf32 data type is faster than standard float32 +torch.backends.cuda.matmul.allow_tf32 = True + +# custom functions # +import utils + + +# In[2]: + + +def classPrecision(logits, y_true, top=1): + """ + Calculate the precision of the top-n predictions. + + Parameters: + logits (torch.Tensor): The output logits from the model (shape: [batch_size, num_classes]). + y_true (torch.Tensor): The ground truth labels (shape: [batch_size]). + top (int): The number of top predictions to consider. + + Returns: + float: The precision percentage of the top-n predictions. + """ + # Apply softmax to get probabilities + probs = F.softmax(logits, dim=1).detach().cpu() + + # Get the top-n predictions + top_n_preds = torch.topk(probs, top, dim=1).indices.detach().cpu() + + # Move y_true to CPU and detach + y_true = y_true.detach().cpu() + + # Check if y_true is in top-n predictions + correct = top_n_preds.eq(y_true.view(-1, 1).expand_as(top_n_preds)) + + # Calculate precision + precision = correct.sum().item() / y_true.size(0) + + return precision * 100 + +# Example usage: +logits = torch.randn(8, 41) # Example logits tensor +y_true = torch.randint(0, 41, (8,)) # Example ground truth labels + +top_n_precision = classPrecision(logits, y_true, top=1) +print(f"Top-1 Precision: {top_n_precision:.2f}%") + + +# In[3]: + + +### Multi-GPU config ### +local_rank = os.getenv('RANK') +if local_rank is None: + local_rank = 0 +else: + local_rank = int(local_rank) +print("LOCAL RANK ", local_rank) + +data_type = torch.float16 # change depending on your mixed_precision +num_devices = torch.cuda.device_count() +if num_devices==0: num_devices = 1 + +# First use "accelerate config" in terminal and setup using deepspeed stage 2 with CPU offloading! +accelerator = Accelerator(split_batches=False, mixed_precision="fp16") +if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job + global_batch_size = batch_size = 16 +else: + global_batch_size = os.environ["GLOBAL_BATCH_SIZE"] + batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices + + +# In[4]: + + +print("PID of this process =",os.getpid()) +device = accelerator.device +print("device:",device) +world_size = accelerator.state.num_processes +distributed = not accelerator.state.distributed_type == 'NO' +num_devices = torch.cuda.device_count() +if num_devices==0 or not distributed: num_devices = 1 +num_workers = num_devices +print(accelerator.state) + +print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type) +print = accelerator.print # only print if local_rank=0 + + +# In[5]: + + +# if running this interactively, can specify jupyter_args here for argparser to use +if utils.is_interactive(): + model_name = "semantic_cluster_0.1" + print("model_name:", model_name) + + # global_batch_size and batch_size should already be defined in the 2nd cell block + jupyter_args = f"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \ + --cache_dir=/weka/proj-medarc/shared/cache \ + --model_name={model_name} \ + --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \ + --hidden_dim=1024 --clip_scale=1. \ + --no-blurry_recon --blur_scale=.5 \ + --use_prior --prior_scale=30 \ + --n_blocks=4 --max_lr=1e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \ + --ckpt_interval=999 --no-ckpt_saving --wandb_log" + # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow + + print(jupyter_args) + jupyter_args = jupyter_args.split() + + from IPython.display import clear_output # function to clear print outputs in cell + get_ipython().run_line_magic('load_ext', 'autoreload') + # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions + get_ipython().run_line_magic('autoreload', '2') + + +# In[6]: + + +parser = argparse.ArgumentParser(description="Model Training Configuration") +parser.add_argument( + "--model_name", type=str, default="testing2", + help="name of model, used for ckpt saving and wandb logging (if enabled)", +) +parser.add_argument( + "--data_path", type=str, default=os.getcwd(), + help="Path to where NSD data is stored / where to download it to", +) +parser.add_argument( + "--cache_dir", type=str, default=os.getcwd(), + help="Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.", +) +parser.add_argument( + "--subj",type=int, default=1, choices=[1,2,3,4,5,6,7,8], + help="Validate on which subject?", +) +parser.add_argument( + "--multisubject_ckpt", type=str, default=None, + help="Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.", +) +parser.add_argument( + "--num_sessions", type=int, default=1, + help="Number of training sessions to include", +) +parser.add_argument( + "--use_prior",action=argparse.BooleanOptionalAction,default=True, + help="whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)", +) +parser.add_argument( + "--batch_size", type=int, default=16, + help="Batch size can be increased by 10x if only training retreival submodule and not diffusion prior", +) +parser.add_argument( + "--wandb_log",action=argparse.BooleanOptionalAction,default=False, + help="whether to log to wandb", +) +parser.add_argument( + "--wandb_project",type=str,default="stability", + help="wandb project name", +) +parser.add_argument( + "--mixup_pct",type=float,default=.33, + help="proportion of way through training when to switch from BiMixCo to SoftCLIP", +) +parser.add_argument( + "--blurry_recon",action=argparse.BooleanOptionalAction,default=True, + help="whether to output blurry reconstructions", +) +parser.add_argument( + "--blur_scale",type=float,default=.5, + help="multiply loss from blurry recons by this number", +) +parser.add_argument( + "--clip_scale",type=float,default=1., + help="multiply contrastive loss by this number", +) +parser.add_argument( + "--prior_scale",type=float,default=30, + help="multiply diffusion prior loss by this", +) +parser.add_argument( + "--use_image_aug",action=argparse.BooleanOptionalAction,default=False, + help="whether to use image augmentation", +) +parser.add_argument( + "--num_epochs",type=int,default=150, + help="number of epochs of training", +) +parser.add_argument( + "--multi_subject",action=argparse.BooleanOptionalAction,default=False, +) +parser.add_argument( + "--new_test",action=argparse.BooleanOptionalAction,default=True, +) +parser.add_argument( + "--n_blocks",type=int,default=4, +) +parser.add_argument( + "--hidden_dim",type=int,default=1024, +) +parser.add_argument( + "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'], +) +parser.add_argument( + "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True, +) +parser.add_argument( + "--ckpt_interval",type=int,default=5, + help="save backup ckpt and reconstruct every x epochs", +) +parser.add_argument( + "--seed",type=int,default=42, +) +parser.add_argument( + "--max_lr",type=float,default=3e-5, +) + +if utils.is_interactive(): + args = parser.parse_args(jupyter_args) +else: + args = parser.parse_args() + +# create global variables without the args prefix +for attribute_name in vars(args).keys(): + globals()[attribute_name] = getattr(args, attribute_name) + +# seed all random functions +utils.seed_everything(seed) + +outdir = os.path.abspath(f'../train_logs/{model_name}') +if not os.path.exists(outdir) and ckpt_saving: + os.makedirs(outdir,exist_ok=True) + +if use_image_aug or blurry_recon: + import kornia + from kornia.augmentation.container import AugmentationSequential +if use_image_aug: + img_augment = AugmentationSequential( + kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3), + same_on_batch=False, + data_keys=["input"], + ) + +if multi_subject: + subj_list = np.arange(1,9) + subj_list = subj_list[subj_list != subj] +else: + subj_list = [subj] + +print("subj_list", subj_list, "num_sessions", num_sessions) + + +# In[7]: + + +def my_split_by_node(urls): return urls +num_voxels_list = [] + +if multi_subject: + nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30]) + num_samples_per_epoch = (750*40) // num_devices +else: + num_samples_per_epoch = (750*num_sessions) // num_devices + +print("dividing batch size by subj_list, which will then be concatenated across subj during training...") +batch_size = batch_size // len(subj_list) + +num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list)) + +print("batch_size =", batch_size, "num_iterations_per_epoch =",num_iterations_per_epoch, "num_samples_per_epoch =",num_samples_per_epoch) + + +# In[8]: + + +train_data = {} +train_dl = {} +num_voxels = {} +voxels = {} +for s in subj_list: + print(f"Training with {num_sessions} sessions") + if multi_subject: + train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{nsessions_allsubj[s-1]-1}" + "}.tar" + else: + train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{num_sessions-1}" + "}.tar" + print(train_url) + + train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,nodesplitter=my_split_by_node)\ + .shuffle(750, initial=1500, rng=random.Random(42))\ + .decode("torch")\ + .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\ + .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"]) + train_dl[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True) + + f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r') + betas = f['betas'][:] + betas = torch.Tensor(betas).to("cpu").to(data_type) + num_voxels_list.append(betas[0].shape[-1]) + num_voxels[f'subj0{s}'] = betas[0].shape[-1] + voxels[f'subj0{s}'] = betas + print(f"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}") + +print("Loaded all subj train dls and betas!\n") + +# Validate only on one subject +if multi_subject: + subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list +if not new_test: # using old test set from before full dataset released (used in original MindEye paper) + if subj==3: + num_test=2113 + elif subj==4: + num_test=1985 + elif subj==6: + num_test=2113 + elif subj==8: + num_test=1985 + else: + num_test=2770 + test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar" +elif new_test: # using larger test set from after full dataset released + if subj==3: + num_test=2371 + elif subj==4: + num_test=2188 + elif subj==6: + num_test=2371 + elif subj==8: + num_test=2188 + else: + num_test=3000 + test_url = f"{data_path}/wds/subj0{subj}/new_test/" + "0.tar" +print(test_url) +test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\ + .shuffle(750, initial=1500, rng=random.Random(42))\ + .decode("torch")\ + .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\ + .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"]) +test_dl = torch.utils.data.DataLoader(test_data, batch_size=num_test, shuffle=False, drop_last=True, pin_memory=True) +print(f"Loaded test dl for subj{subj}!\n") + + +# In[9]: + + +# Load 73k NSD images +f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r') +images = f['images'] +print("Loaded all 73k possible NSD images to cpu!", images.shape) + + +# In[10]: + + +clip_img_embedder = FrozenOpenCLIPImageEmbedder( + arch="ViT-bigG-14", + version="laion2b_s39b_b160k", + output_tokens=True, + only_tokens=True, +) +clip_img_embedder.to(device) + +clip_seq_dim = 256 +clip_emb_dim = 1664 + + +# In[11]: + + +if blurry_recon: + from diffusers import AutoencoderKL + autoenc = AutoencoderKL( + down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'], + up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'], + block_out_channels=[128, 256, 512, 512], + layers_per_block=2, + sample_size=256, + ) + ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth') + autoenc.load_state_dict(ckpt) + + autoenc.eval() + autoenc.requires_grad_(False) + autoenc.to(device) + utils.count_params(autoenc) + + from autoencoder.convnext import ConvnextXL + cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth') + cnx.requires_grad_(False) + cnx.eval() + cnx.to(device) + + mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1) + std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1) + + blur_augs = AugmentationSequential( + kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8), + kornia.augmentation.RandomGrayscale(p=0.1), + kornia.augmentation.RandomSolarize(p=0.1), + kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0), + data_keys=["input"], + ) + + +# In[12]: + + +class MindEyeModule(nn.Module): + def __init__(self): + super(MindEyeModule, self).__init__() + def forward(self, x): + return x + +model = MindEyeModule() +model + + +# In[13]: + + +class RidgeRegression(torch.nn.Module): + # make sure to add weight_decay when initializing optimizer to enable regularization + def __init__(self, input_sizes, out_features): + super(RidgeRegression, self).__init__() + self.out_features = out_features + self.linears = torch.nn.ModuleList([ + torch.nn.Linear(input_size, out_features) for input_size in input_sizes + ]) + def forward(self, x, subj_idx): + out = self.linears[subj_idx](x[:,0]).unsqueeze(1) + return out + +class IndividRidgeRegression(torch.nn.Module): + def __init__(self, input_size, out_features): + super(IndividRidgeRegression, self).__init__() + self.out_features = out_features + self.linear = torch.nn.Linear(input_size, out_features) + def forward(self, x): + out = self.linear(x) + return out + +model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim) +utils.count_params(model.ridge) +utils.count_params(model) + +# test on subject 1 with fake data +b = torch.randn((2,1,num_voxels_list[0])) +print(b.shape, model.ridge(b,0).shape) + + +# In[14]: + + +from models import BrainNetwork +model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks, + clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, + blurry_recon=blurry_recon, clip_scale=clip_scale) +utils.count_params(model.backbone) +utils.count_params(model) + +# test that the model works on some fake data +b = torch.randn((2,1,hidden_dim)) +print("b.shape",b.shape) + +backbone_, clip_, blur_ = model.backbone(b) +print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape) + + +# In[15]: + + +if use_prior: + from models import * + + # setup diffusion prior network + out_dim = clip_emb_dim + depth = 6 + dim_head = 52 + heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim + timesteps = 100 + + prior_network = PriorNetwork( + dim=out_dim, + depth=depth, + dim_head=dim_head, + heads=heads, + causal=False, + num_tokens = clip_seq_dim, + learned_query_mode="pos_emb" + ) + + model.diffusion_prior = BrainDiffusionPrior( + net=prior_network, + image_embed_dim=out_dim, + condition_on_text_encodings=False, + timesteps=timesteps, + cond_drop_prob=0.2, + image_embed_scale=None, + ) + + utils.count_params(model.diffusion_prior) + utils.count_params(model) + + +# In[16]: + + +path_semantic_names = "/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy" +path_semantic_cluster = "/weka/proj-fmri/ckadirt/MindEyeV2/src/COCO_73k_semantic_cluster.npy" +semantic_cluster_names = np.load(path_semantic_names) +semantic_cluster = np.load(path_semantic_cluster) +possible_semantic_clusters = np.unique(semantic_cluster) + +# one-hot encode semantic clusters +# move possible_semantic_clusters to numbers and create a dictionary +semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)} +semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters))) +for i, cluster in enumerate(semantic_cluster): + semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1 + + +print("semantic_cluster_onehot.shape", semantic_cluster_onehot.shape) + +num_seman_clusters = len(np.unique(semantic_cluster)) +print("num_seman_clusters", num_seman_clusters) + + +# In[17]: + + +# plot some images next to their semantic cluster +fig, ax = plt.subplots(1, 5, figsize=(20, 4)) +for i in range(5): + # covert numpy array images to float32 + image_index = torch.randint(0, len(images), (1,)).item() + print(image_index) + ax[i].imshow(images[image_index].transpose(1,2,0).astype(np.float32)) + ax[i].set_title(semantic_cluster[image_index]) + ax[i].axis("off") +plt.show() + + +# In[18]: + + +model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters) +utils.count_params(model.RRClassifier) +utils.count_params(model) + + +# In[19]: + + +no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] + +opt_grouped_parameters = [ + {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2}, + {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2}, + {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}, + {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1}, +] +# if use_prior: +# opt_grouped_parameters.extend([ +# {'params': [p for n, p in model.diffusion_prior.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2}, +# {'params': [p for n, p in model.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} +# ]) +# opt_grouped_parameters.extend([ +# +# ]) + +optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr) + +if lr_scheduler_type == 'linear': + lr_scheduler = torch.optim.lr_scheduler.LinearLR( + optimizer, + total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)), + last_epoch=-1 + ) +elif lr_scheduler_type == 'cycle': + total_steps=int(np.floor(num_epochs*num_iterations_per_epoch)) + print("total_steps", total_steps) + lr_scheduler = torch.optim.lr_scheduler.OneCycleLR( + optimizer, + max_lr=max_lr, + total_steps=total_steps, + final_div_factor=1000, + last_epoch=-1, pct_start=2/num_epochs + ) + +def save_ckpt(tag): + ckpt_path = outdir+f'/{tag}.pth' + if accelerator.is_main_process: + unwrapped_model = accelerator.unwrap_model(model) + torch.save({ + 'epoch': epoch, + 'model_state_dict': unwrapped_model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'lr_scheduler': lr_scheduler.state_dict(), + 'train_losses': losses, + 'test_losses': test_losses, + 'lrs': lrs, + }, ckpt_path) + print(f"\n---saved {outdir}/{tag} ckpt!---\n") + +def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): + print(f"\n---loading {outdir}/{tag}.pth ckpt---\n") + checkpoint = torch.load(outdir+'/last.pth', map_location='cpu') + state_dict = checkpoint['model_state_dict'] + if multisubj_loading: # remove incompatible ridge layer that will otherwise error + state_dict.pop('ridge.linears.0.weight',None) + model.load_state_dict(state_dict, strict=strict) + if load_epoch: + globals()["epoch"] = checkpoint['epoch'] + print("Epoch",epoch) + if load_optimizer: + optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + if load_lr: + lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) + del checkpoint + +print("\nDone with model preparations!") +num_params = utils.count_params(model) + + +# In[20]: + + +max_lr + + +# In[21]: + + +if local_rank==0 and wandb_log: # only use main process for wandb logging + import wandb + wandb_project = 'mindeye_semantic_cluster_0.2' + print(f"wandb {wandb_project} run {model_name}") + # need to configure wandb beforehand in terminal with "wandb init"! + wandb_config = { + "model_name": model_name, + "global_batch_size": global_batch_size, + "batch_size": batch_size, + "num_epochs": num_epochs, + "num_sessions": num_sessions, + "num_params": num_params, + "clip_scale": clip_scale, + "prior_scale": prior_scale, + "blur_scale": blur_scale, + "use_image_aug": use_image_aug, + "max_lr": max_lr, + "mixup_pct": mixup_pct, + "num_samples_per_epoch": num_samples_per_epoch, + "num_test": num_test, + "ckpt_interval": ckpt_interval, + "ckpt_saving": ckpt_saving, + "seed": seed, + "distributed": distributed, + "num_devices": num_devices, + "world_size": world_size, + "train_url": train_url, + "test_url": test_url, + } + print("wandb_config:\n",wandb_config) + print("wandb_id:",model_name) + wandb.login(host='https://stability.wandb.io') + wandb.init( + id=model_name, + project=wandb_project, + name=model_name, + config=wandb_config, + resume="allow", + ) +else: + wandb_log = False + + +# In[22]: + + +epoch = 0 +losses, test_losses, lrs = [], [], [] +best_test_loss = 1e9 +torch.cuda.empty_cache() + + +# In[23]: + + +# load multisubject stage1 ckpt if set +if multisubject_ckpt is not None: + load_ckpt("last",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True) + + +# In[24]: + + +train_dls = [train_dl[f'subj0{s}'] for s in subj_list] + +model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot) +# leaving out test_dl since we will only have local_rank 0 device do evals + + +# In[25]: + + +def plot_semantic_clusters(images, indexes, semantic_cluster, semantic_cluster_dict): + # check images and indexes are the same length + assert len(images) == len(indexes) + + # if images are tensors, convert them to numpy arrays and move them to cpu + if isinstance(images, torch.Tensor): + images = images.cpu().numpy() + + fig, ax = plt.subplots(1, len(images), figsize=(20, 4)) + for i, index in enumerate(indexes): + # covert numpy array images to float32 + ax[i].imshow(images[i].transpose(1,2,0).astype(np.float32)) + # search the key in the dictionary based on the index as value + name = {i for i in semantic_cluster_dict if semantic_cluster_dict[i] == index} + ax[i].set_title(name) + ax[i].axis("off") + plt.show() + + +# In[26]: + + +#RTTT + + +# In[30]: + + +print(f"{model_name} starting with epoch {epoch} / {num_epochs}") +progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0)) +test_image, test_voxel = None, None +mse = nn.MSELoss() +l1 = nn.L1Loss() +soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs)) + +for epoch in progress_bar: + model.train() + + fwd_percent_correct = 0. + bwd_percent_correct = 0. + test_fwd_percent_correct = 0. + test_bwd_percent_correct = 0. + + recon_cossim = 0. + test_recon_cossim = 0. + recon_mse = 0. + test_recon_mse = 0. + + loss_clip_total = 0. + loss_blurry_total = 0. + loss_blurry_cont_total = 0. + test_loss_clip_total = 0. + + loss_prior_total = 0. + test_loss_prior_total = 0. + + blurry_pixcorr = 0. + test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1 + + loss_RR_total = 0. + test_loss_RR_total = 0. + + class_precisions_1 = 0 + test_class_precisions_1 = 0 + + class_precisions_5 = 0 + test_class_precisions_5 = 0 + + class_precisions_10 = 0 + test_class_precisions_10 = 0 + + # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch) + voxel_iters = {} # empty dict because diff subjects have differing # of voxels + image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float() + annot_iters = {} + perm_iters, betas_iters, select_iters = {}, {}, {} + images_indexes = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list)).long() + for s, train_dl in enumerate(train_dls): + with torch.cuda.amp.autocast(dtype=data_type): + iter = -1 + for behav0, past_behav0, future_behav0, old_behav0 in train_dl: + # Load images to cpu from hdf5 (requires sorted indexing) + image_idx = behav0[:,0,0].cpu().long().numpy() + image0, image_sorted_idx = np.unique(image_idx, return_index=True) + + if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing + continue + iter += 1 + images_indexes[iter, s*batch_size:s*batch_size+batch_size] = torch.Tensor(image0) + image0 = torch.tensor(images[image0], dtype=data_type) + #print(image0.shape) + image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0 + + # Load voxels for current batch, matching above indexing + voxel_idx = behav0[:,0,5].cpu().long().numpy() + voxel_sorted_idx = voxel_idx[image_sorted_idx] + voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx] + voxel0 = torch.Tensor(voxel0).unsqueeze(1) + + if epoch < int(mixup_pct * num_epochs): + voxel0, perm, betas, select = utils.mixco(voxel0) + perm_iters[f"subj0{subj_list[s]}_iter{iter}"] = perm + betas_iters[f"subj0{subj_list[s]}_iter{iter}"] = betas + select_iters[f"subj0{subj_list[s]}_iter{iter}"] = select + + voxel_iters[f"subj0{subj_list[s]}_iter{iter}"] = voxel0 + + if iter >= num_iterations_per_epoch-1: + break + + # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each + for train_i in range(num_iterations_per_epoch): + with torch.cuda.amp.autocast(dtype=data_type): + + optimizer.zero_grad() + loss=0. + + voxel_list = [voxel_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list] + image = image_iters[train_i].detach() + image = image.to(device) + + if use_image_aug: + image = img_augment(image) + + clip_target = clip_img_embedder(image) + assert not torch.any(torch.isnan(clip_target)) + + if epoch < int(mixup_pct * num_epochs): + perm_list = [perm_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list] + perm = torch.cat(perm_list, dim=0) + betas_list = [betas_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list] + betas = torch.cat(betas_list, dim=0) + select_list = [select_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list] + select = torch.cat(select_list, dim=0) + + voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)] + voxel_ridge = torch.cat(voxel_ridge_list, dim=0) + + backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge) + logits = model.RRClassifier(backbone.flatten(1)) + #print(semantic_cluster[images_indexes[f"subj0{s}_iter{train_i}"]) + labels = semantic_cluster[images_indexes[train_i].type(torch.LongTensor).tolist()] + indexes = torch.Tensor([semantic_cluster_dict[i] for i in labels]).type(torch.LongTensor) + loss_SM = nn.functional.cross_entropy(logits, indexes.to(logits.device)) + + # plot_semantic_clusters(image[0:5], indexes[0:5].to('cpu'), semantic_cluster, semantic_cluster_dict) + + loss+= loss_SM + loss_RR_total += loss_SM.item() + + if (torch.rand(1) < 0.03).item(): + print("loss_SM", loss_SM.item()) + + if clip_scale>0: + clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1) + clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1) + + if use_prior: + loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target) + loss_prior_total += loss_prior.item() + loss_prior *= prior_scale + loss += loss_prior + + recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item() + recon_mse += mse(prior_out, clip_target).item() + + if clip_scale>0: + if epoch < int(mixup_pct * num_epochs): + loss_clip = utils.mixco_nce( + clip_voxels_norm, + clip_target_norm, + temp=.006, + perm=perm, betas=betas, select=select) + else: + epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)] + loss_clip = utils.soft_clip_loss( + clip_voxels_norm, + clip_target_norm, + temp=epoch_temp) + + loss_clip_total += loss_clip.item() + loss_clip *= clip_scale + loss += loss_clip + + if blurry_recon: + image_enc_pred, transformer_feats = blurry_image_enc_ + + image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215 + loss_blurry = l1(image_enc_pred, image_enc) + loss_blurry_total += loss_blurry.item() + + if epoch < int(mixup_pct * num_epochs): + image_enc_shuf = image_enc[perm] + betas_shape = [-1] + [1]*(len(image_enc.shape)-1) + image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \ + image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape) + + image_norm = (image - mean)/std + image_aug = (blur_augs(image) - mean)/std + _, cnx_embeds = cnx(image_norm) + _, cnx_aug_embeds = cnx(image_aug) + + cont_loss = utils.soft_cont_loss( + nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1), + nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1), + nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1), + temp=0.2) + loss_blurry_cont_total += cont_loss.item() + + loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215 + + if clip_scale>0: + # forward and backward top 1 accuracy + labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) + fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item() + bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item() + + class_precisions_1 += classPrecision(logits, indexes) + class_precisions_5 += classPrecision(logits, indexes, 5) + class_precisions_10 += classPrecision(logits, indexes, 10) + + if blurry_recon: + with torch.no_grad(): + # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode() + random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False) + blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1) + pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images) + blurry_pixcorr += pixcorr.item() + + utils.check_loss(loss) + accelerator.backward(loss) + optimizer.step() + + losses.append(loss.item()) + lrs.append(optimizer.param_groups[0]['lr']) + + if lr_scheduler_type is not None: + lr_scheduler.step() + + model.eval() + if local_rank==0: + with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): + for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl): + # all test samples should be loaded per batch such that test_i should never exceed 0 + assert len(behav) == num_test + + ## Average same-image repeats ## + if test_image is None: + test_coco_in = [] + voxel = voxels[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1) + + image = behav[:,0,0].cpu().long() + + unique_image, sort_indices = torch.unique(image, return_inverse=True) + for im in unique_image: + locs = torch.where(im == image)[0] + if len(locs)==1: + locs = locs.repeat(3) + elif len(locs)==2: + locs = locs.repeat(2)[:3] + assert len(locs)==3 + if test_image is None: + test_image = torch.Tensor(images[im][None]) + test_voxel = voxel[locs][None] + else: + test_image = torch.vstack((test_image, torch.Tensor(images[im][None]))) + test_voxel = torch.vstack((test_voxel, voxel[locs][None])) + test_coco_in.append(int(im)) + + loss=0. + #print(test_coco_in, len(test_coco_in)) + test_indices = torch.arange(len(test_voxel))[:300] + voxel = test_voxel[test_indices].to(device) + image = test_image[test_indices].to(device) + assert len(image) == 300 + + clip_target = clip_img_embedder(image.float()) + + for rep in range(3): + voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list + backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge) + logits0 = model.RRClassifier(backbone0.flatten(1)) + + if rep==0: + clip_voxels = clip_voxels0 + backbone = backbone0 + logits = logits0 + else: + clip_voxels += clip_voxels0 + backbone += backbone0 + logits += logits0 + + clip_voxels /= 3 + backbone /= 3 + logits /= 3 + + + #print(semantic_cluster[images_indexes[f"subj0{s}_iter{train_i}"]) + #print(test_indices) + labels = semantic_cluster[torch.Tensor(test_coco_in)[test_indices].long().tolist()] + indexes = torch.Tensor([semantic_cluster_dict[i] for i in labels]).type(torch.LongTensor) + loss_SM = nn.functional.cross_entropy(logits, indexes.to(logits.device)) + #plot_semantic_clusters(image[0:5], indexes[0:5].to('cpu'), semantic_cluster, semantic_cluster_dict) + + loss+= loss_SM + test_loss_RR_total += loss_SM.item() + if clip_scale>0: + clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1) + clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1) + + # for some evals, only doing a subset of the samples per batch because of computational cost + random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False) + + if use_prior: + loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps]) + test_loss_prior_total += loss_prior.item() + loss_prior *= prior_scale + loss += loss_prior + + if clip_scale>0: + loss_clip = utils.soft_clip_loss( + clip_voxels_norm, + clip_target_norm, + temp=.006) + + test_loss_clip_total += loss_clip.item() + loss_clip = loss_clip * clip_scale + loss += loss_clip + + if blurry_recon: + image_enc_pred, _ = blurry_image_enc_ + blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1) + pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images) + test_blurry_pixcorr += pixcorr.item() + + if clip_scale>0: + # forward and backward top 1 accuracy + labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) + test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item() + test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item() + + test_class_precisions_1 += classPrecision(logits, indexes) + test_class_precisions_5 += classPrecision(logits, indexes, 5) + test_class_precisions_10 += classPrecision(logits, indexes, 10) + + utils.check_loss(loss) + test_losses.append(loss.item()) + + + assert (test_i+1) == 1 + logs = {"train/loss": np.mean(losses[-(train_i+1):]), + "test/loss": np.mean(test_losses[-(test_i+1):]), + #"train/lr": lrs[-1], + "train/num_steps": len(losses), + "test/num_steps": len(test_losses), + "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1), + "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1), + "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1), + "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1), + "train/loss_clip_total": loss_clip_total / (train_i + 1), + "train/loss_blurry_total": loss_blurry_total / (train_i + 1), + "train/loss_blurry_cont_total": loss_blurry_cont_total / (train_i + 1), + "test/loss_clip_total": test_loss_clip_total / (test_i + 1), + "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1), + "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1), + "train/recon_cossim": recon_cossim / (train_i + 1), + "test/recon_cossim": test_recon_cossim / (test_i + 1), + "train/recon_mse": recon_mse / (train_i + 1), + "test/recon_mse": test_recon_mse / (test_i + 1), + "train/loss_prior": loss_prior_total / (train_i + 1), + "test/loss_prior": test_loss_prior_total / (test_i + 1), + "train/loss_RR": loss_RR_total / (train_i + 1), + "test/loss_RR": test_loss_RR_total / (test_i + 1), + "train/class_precisions_1": class_precisions_1 / (train_i + 1), + "test/class_precisions_1": test_class_precisions_1 / (test_i + 1), + "train/class_precisions_5": class_precisions_5 / (train_i + 1), + "test/class_precisions_5": test_class_precisions_5 / (test_i + 1), + "train/class_precisions_10": class_precisions_10 / (train_i + 1), + "test/class_precisions_10": test_class_precisions_10 / (test_i + 1), + } + + # if finished training, save jpg recons if they exist + if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0): + if blurry_recon: + image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215 + # transform blurry recon latents to images and plot it + fig, axes = plt.subplots(1, 8, figsize=(10, 4)) + jj=-1 + for j in [0,1,2,3]: + jj+=1 + axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1))) + axes[jj].axis('off') + jj+=1 + axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1))) + axes[jj].axis('off') + + if wandb_log: + logs[f"test/blur_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}") + plt.close() + else: + plt.show() + + progress_bar.set_postfix(**logs) + + if wandb_log: wandb.log(logs) + + # Save model checkpoint and reconstruct + if (ckpt_saving) and (epoch % ckpt_interval == 0): + save_ckpt(f'last') + + # wait for other GPUs to catch up if needed + accelerator.wait_for_everyone() + torch.cuda.empty_cache() + +print("\n===Finished!===\n") +if ckpt_saving: + save_ckpt(f'last') + + +# In[29]: + + +test_coco_in + diff --git a/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/config.yaml b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6bdc3cb68319beaf195592c3d9539e32140d2b7a --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/config.yaml @@ -0,0 +1,115 @@ +wandb_version: 1 + +model_name: + desc: null + value: semantic_cluster_1.2_average_after_wd-2_no_prior_multi +global_batch_size: + desc: null + value: '21' +batch_size: + desc: null + value: 3 +num_epochs: + desc: null + value: 150 +num_sessions: + desc: null + value: 40 +num_params: + desc: null + value: 573919937 +clip_scale: + desc: null + value: 1.0 +prior_scale: + desc: null + value: 30.0 +blur_scale: + desc: null + value: 0.5 +use_image_aug: + desc: null + value: false +max_lr: + desc: null + value: 3.0e-05 +mixup_pct: + desc: null + value: 0.1 +num_samples_per_epoch: + desc: null + value: 30000 +num_test: + desc: null + value: 3000 +ckpt_interval: + desc: null + value: 999 +ckpt_saving: + desc: null + value: true +seed: + desc: null + value: 42 +distributed: + desc: null + value: false +num_devices: + desc: null + value: 1 +world_size: + desc: null + value: 1 +train_url: + desc: null + value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj08/train/{0..29}.tar +test_url: + desc: null + value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj02/new_test/0.tar +_wandb: + desc: null + value: + code_path: code/src/Untitled1.py + python_version: 3.11.9 + cli_version: 0.17.1 + framework: huggingface + huggingface_version: 4.37.2 + is_jupyter_run: false + is_kaggle_kernel: false + start_time: 1720572995 + t: + 1: + - 1 + - 9 + - 11 + - 41 + - 49 + - 55 + - 63 + - 71 + - 79 + - 83 + - 103 + 2: + - 1 + - 9 + - 11 + - 41 + - 49 + - 55 + - 63 + - 71 + - 79 + - 83 + - 103 + 3: + - 13 + - 14 + - 16 + - 23 + 4: 3.11.9 + 5: 0.17.1 + 6: 4.37.2 + 8: + - 5 + 13: linux-x86_64 diff --git a/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/diff.patch b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/diff.patch new file mode 100644 index 0000000000000000000000000000000000000000..5d22c2032a18410574e893054181202921670d1a --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/diff.patch @@ -0,0 +1,157 @@ +diff --git a/.gitignore b/.gitignore +old mode 100644 +new mode 100755 +diff --git a/LICENSE b/LICENSE +old mode 100644 +new mode 100755 +diff --git a/README.md b/README.md +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_1sess_24bs_all_blurryrecons_1000recons.png b/figs/final_subj01_pretrained_1sess_24bs_all_blurryrecons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_1sess_24bs_all_enhancedrecons_1000recons.png b/figs/final_subj01_pretrained_1sess_24bs_all_enhancedrecons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_1sess_24bs_all_recons_1000recons.png b/figs/final_subj01_pretrained_1sess_24bs_all_recons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_2sess_24bs_all_enhancedrecons_1000recons.png b/figs/final_subj01_pretrained_2sess_24bs_all_enhancedrecons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_375trials_24bs_all_enhancedrecons_1000recons.png b/figs/final_subj01_pretrained_375trials_24bs_all_enhancedrecons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_40sess_24bs_all_blurryrecons_1000recons.png b/figs/final_subj01_pretrained_40sess_24bs_all_blurryrecons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_40sess_24bs_all_enhancedrecons_1000recons.png b/figs/final_subj01_pretrained_40sess_24bs_all_enhancedrecons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_40sess_24bs_all_recons_1000recons.png b/figs/final_subj01_pretrained_40sess_24bs_all_recons_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj01_pretrained_40sess_24bs_umap_viz.png b/figs/final_subj01_pretrained_40sess_24bs_umap_viz.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj02_pretrained_1sess_24bs_all_enhancedrec_1000recons.png b/figs/final_subj02_pretrained_1sess_24bs_all_enhancedrec_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj02_pretrained_40sess_24bs_all_enhancedrec_1000recons.png b/figs/final_subj02_pretrained_40sess_24bs_all_enhancedrec_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj02_pretrained_40sess_24bs_umap_viz.png b/figs/final_subj02_pretrained_40sess_24bs_umap_viz.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj05_pretrained_1sess_24bs_all_enhancedrec_1000recons.png b/figs/final_subj05_pretrained_1sess_24bs_all_enhancedrec_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj05_pretrained_40sess_24bs_all_enhancedrec_1000recons.png b/figs/final_subj05_pretrained_40sess_24bs_all_enhancedrec_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj05_pretrained_40sess_24bs_umap_viz.png b/figs/final_subj05_pretrained_40sess_24bs_umap_viz.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj07_pretrained_1sess_24bs_all_enhancedrec_1000recons.png b/figs/final_subj07_pretrained_1sess_24bs_all_enhancedrec_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj07_pretrained_40sess_24bs_all_enhancedrec_1000recons.png b/figs/final_subj07_pretrained_40sess_24bs_all_enhancedrec_1000recons.png +old mode 100644 +new mode 100755 +diff --git a/figs/final_subj07_pretrained_40sess_24bs_umap_viz.png b/figs/final_subj07_pretrained_40sess_24bs_umap_viz.png +old mode 100644 +new mode 100755 +diff --git a/figs/gt_umap_viz.png b/figs/gt_umap_viz.png +old mode 100644 +new mode 100755 +diff --git a/figs/recon_comparison_small_alt.png b/figs/recon_comparison_small_alt.png +old mode 100644 +new mode 100755 +diff --git a/src/Train.ipynb b/src/Train.ipynb +old mode 100644 +new mode 100755 +index 91922c8..396549d +--- a/src/Train.ipynb ++++ b/src/Train.ipynb +@@ -99,7 +99,7 @@ + "name": "stdout", + "output_type": "stream", + "text": [ +- "PID of this process = 3001293\n", ++ "PID of this process = 3520819\n", + "device: cuda\n", + "Distributed environment: DistributedType.NO\n", + "Num processes: 1\n", +@@ -1362,9 +1362,9 @@ + ], + "metadata": { + "kernelspec": { +- "display_name": "mindeye", ++ "display_name": "Python 3 (ipykernel)", + "language": "python", +- "name": "mindeye" ++ "name": "python3" + }, + "language_info": { + "codemirror_mode": { +diff --git a/src/accel.slurm b/src/accel.slurm +old mode 100644 +new mode 100755 +index 1c1dfbc..f0b347f +--- a/src/accel.slurm ++++ b/src/accel.slurm +@@ -1,7 +1,8 @@ + #!/bin/bash +-#SBATCH --account=topfmri +-#SBATCH --partition=a40x +-#SBATCH --job-name=eye ++#SBATCH --account=fmri ++#SBATCH --qos=normal ++#SBATCH --partition=p5 ++#SBATCH --job-name=eyeRR + #SBATCH --nodes=1 + #SBATCH --gres=gpu:1 + #SBATCH --time=48:00:00 # total run time limit (HH:MM:SS) +@@ -12,11 +13,11 @@ + #SBATCH --exclusive + + # Make sure you activate your fmri environment created from src/setup.sh +-cd /weka/proj-fmri/paulscotti/MindEyeV2/src +-source fmri/bin/activate ++cd /weka/proj-fmri/ckadirt/MindEyeV2/src ++source /admin/home-ckadirt/fmri/bin/activate + + # The following line converts your jupyter notebook into a python script runnable with Slurm +-jupyter nbconvert Train.ipynb --to python ++jupyter nbconvert TrainCluster-Copy2.ipynb --to python + + export NUM_GPUS=1 # Set to equal gres=gpu:#! + export BATCH_SIZE=21 # 21 for multisubject / 24 for singlesubject (orig. paper used 42 for multisubject / 24 for singlesubject) +@@ -26,16 +27,16 @@ export GLOBAL_BATCH_SIZE=$((BATCH_SIZE * NUM_GPUS)) + export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) + export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST") + export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +-export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l) ++export COUNT_NODE=1 #$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l) + echo MASTER_ADDR=${MASTER_ADDR} + echo MASTER_PORT=${MASTER_PORT} + echo WORLD_SIZE=${COUNT_NODE} + + # multisubject pretraining +-model_name="multisubject_excludingsubj01_40sess" ++model_name="rr-test-top" + echo model_name=${model_name} +-accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 Train.py --data_path=/weka/proj-fmri/shared/mindeyev2_dataset --cache_dir=/weka/proj-fmri/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --use_prior --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log +- ++# accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 TrainCluster.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log ++python TrainCluster-Copy2.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --no-multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-5 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log + # singlesubject finetuning + #model_name="finetuned_subj01_40sess" + #echo model_name=${model_name} +diff --git a/src/autoencoder/convnext.py b/src/autoencoder/convnext.py +old mode 100644 +new mode 100755 +diff --git a/src/dataset_creation.ipynb b/src/dataset_creation.ipynb +old mode 100644 +new mode 100755 +diff --git a/src/enhanced_recon_inference.ipynb b/src/enhanced_recon_inf \ No newline at end of file diff --git a/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/output.log b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..a0d30c781b9e2fbd2117923db3c9e5a87cddd8bc --- /dev/null +++ b/MindEyeV2/src/wandb/run-20240710_005635-semantic_cluster_1.2_average_after_wd-2_no_prior_multi/files/output.log @@ -0,0 +1,6735 @@ + +semantic_cluster_1.2_average_after_wd-2_no_prior_multi starting with epoch 0 / 150 + 0%| | 0/150 [00:00