text: import os, sys
  #os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
  os.environ["CUDA_VISIBLE_DEVICES"]="1"
  #from fastai.conv_learner import *
  from fastai.conv_learner import RandomDihedral, RandomLighting, TfmType, RandomRotate, CropType, resnet34, cv2, A, ImageData, tfms_from_stats, nn, F, model_meta, cut_model, torch, to_gpu, apply_init, model_features, num_features
  from fastai.conv_learner import AdaptiveConcatPool2d, Flatten, kaiming_normal, children, split_by_idxs, Learner, accuracy_thresh, ConvnetBuilder, FP16, optim, predict_to_bcolz, bcolz
  from Christof.static import name_label_dict, channel_avr, channel_std, lb_prob
  #from fastai.dataset import *
  from fastai.dataset import FilesDataset, ImageClassifierData
  import pandas as pd
  import numpy as np
  import os
  from sklearn.model_selection import train_test_split
  from sklearn.metrics import f1_score
  import scipy.optimize as opt
  
  PATH = './'
  TRAIN = 'Christof/assets/train/'
  TEST = 'Christof/assets/test/'
  LABELS = 'Christof/assets/train.csv'
  SAMPLE = 'Christof/assets/sample_submission.csv'
  
  MODEL_PATH = 'Christof/models/ResNet34/fastai/3_10/'
  
  nw = 2   #number of workers for data loader
  arch = resnet34 #specify target architecture
  
  train_names = list({f[:36] for f in os.listdir(TRAIN)})
  test_names = list({f[:36] for f in os.listdir(TEST)})
  
  
  import pickle
  with open('Russ/folds2.pkl', 'rb') as f:
      folds = pickle.load(f)
  
  def get_fold(fold_id,folds):
      tr_n = list({fn for fn in folds if folds[fn]!=fold_id})
      val_n = list({fn for fn in folds if folds[fn] == fold_id})
      return tr_n, val_n
  
  #tr_n, val_n = train_test_split(train_names, test_size=0.1, random_state=42)
  
  
  
  def open_rgby(path,id): #a function that reads RGBY image
      colors = ['red','green','blue']
      flags = cv2.IMREAD_GRAYSCALE
      img = [cv2.imread(os.path.join(path, id+'_'+color+'.png'), flags).astype(np.float32)/255
             for color in colors]
      return np.stack(img, axis=-1)
  
  
  class pdFilesDataset(FilesDataset):
      def __init__(self, fnames, path, transform):
          self.labels = pd.read_csv(LABELS).set_index('Id')
          self.labels['Target'] = [[int(i) for i in s.split()] for s in self.labels['Target']]
          super().__init__(fnames, transform, path)
  
      def get_x(self, i):
          img = open_rgby(self.path, self.fnames[i])
          if self.sz == 512:
              return img
          else:
              return cv2.resize(img, (self.sz, self.sz), cv2.INTER_AREA)
  
      def get_y(self, i):
          if (self.path == TEST):
              return np.zeros(len(name_label_dict), dtype=np.int)
          else:
              labels = self.labels.loc[self.fnames[i]]['Target']
              return np.eye(len(name_label_dict), dtype=np.float)[labels].sum(axis=0)
  
      @property
      def is_multi(self):
          return True
  
      @property
      def is_reg(self):
          return True
  
      # this flag is set to remove the output sigmoid that allows log(sigmoid) optimization
      # of the numerical stability of the loss function
  
      def get_c(self):
          return len(name_label_dict)  # number of classes
  
  
  def get_data(sz,bs, tr_n,val_n):
      #data augmentation
      aug_tfms = [RandomRotate(30, tfm_y=TfmType.NO),
                  RandomDihedral(tfm_y=TfmType.NO),
                  RandomLighting(0.05, 0.05, tfm_y=TfmType.NO)]
      #mean and std in of each channel in the train set
      stats = A([0.08069, 0.05258, 0.05487], [0.13704, 0.10145, 0.15313])
      tfms = tfms_from_stats(stats, sz, crop_type=CropType.NO, tfm_y=TfmType.NO,
                  aug_tfms=aug_tfms)
      ds = ImageData.get_ds(pdFilesDataset, (tr_n[:-(len(tr_n)%bs)],TRAIN),(val_n,TRAIN), tfms, test=(test_names,TEST))
      md = ImageData(PATH, ds, bs, num_workers=nw, classes=None)
      return md
  
  #bs = 16
  #sz = 256
  #md = get_data(sz,bs)
  
  #x,y = next(iter(md.trn_dl))
  #print(x.shape, y.shape)
  
  print('averages:')
  print(channel_avr,channel_std)
  
  
  def display_imgs(x):
      columns = 4
      bs = x.shape[0]
      rows = min((bs + 3) // 4, 4)
      fig = plt.figure(figsize=(columns * 4, rows * 4))
      for i in range(rows):
          for j in range(columns):
              idx = i + j * columns
              fig.add_subplot(rows, columns, idx + 1)
              plt.axis('off')
              plt.imshow((x[idx, :, :, :3] * 255).astype(np.int))
      plt.show()
  
  
  #display_imgs(np.asarray(md.trn_ds.denorm(x)))
  
  
  # define loss
  
  class FocalLoss(nn.Module):
      def __init__(self, gamma=2):
          super().__init__()
          self.gamma = gamma
  
      def forward(self, input, target):
          if not (target.size() == input.size()):
              raise ValueError("Target size ({}) must be the same as input size ({})"
                               .format(target.size(), input.size()))
  
          max_val = (-input).clamp(min=0)
          loss = input - input * target + max_val + \
                 ((-max_val).exp() + (-input - max_val).exp()).log()
  
          invprobs = F.logsigmoid(-input * (target * 2.0 - 1.0))
          loss = (invprobs * self.gamma).exp() * loss
  
          return loss.sum(dim=1).mean()
  
  def acc(preds,targs,th=0.0):
      preds = (preds > th).int()
      targs = targs.int()
      return (preds==targs).float().mean()
  
  
  class ConvnetBuilder_custom():
      def __init__(self, f, c, is_multi, is_reg, ps=None, xtra_fc=None, xtra_cut=0,
                   custom_head=None, pretrained=True):
          self.f, self.c, self.is_multi, self.is_reg, self.xtra_cut = f, c, is_multi, is_reg, xtra_cut
          if xtra_fc is None: xtra_fc = [512]
          if ps is None: ps = [0.25] * len(xtra_fc) + [0.5]
          self.ps, self.xtra_fc = ps, xtra_fc
  
          if f in model_meta:
              cut, self.lr_cut = model_meta[f]
          else:
              cut, self.lr_cut = 0, 0
          cut -= xtra_cut
          layers = cut_model(f(pretrained), cut)
  
          # replace first convolutional layer by 4->64 while keeping corresponding weights
          # and initializing new weights with zeros
          #w = layers[0].weight
          #layers[0] = nn.Conv2d(4, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
          #layers[0].weight = torch.nn.Parameter(torch.cat((w, torch.zeros(64, 1, 7, 7)), dim=1))
  
  
          self.nf = model_features[f] if f in model_features else (num_features(layers) * 2)
          if not custom_head: layers += [AdaptiveConcatPool2d(), Flatten()]
          self.top_model = nn.Sequential(*layers)
  
          n_fc = len(self.xtra_fc) + 1
          if not isinstance(self.ps, list): self.ps = [self.ps] * n_fc
  
          if custom_head:
              fc_layers = [custom_head]
          else:
              fc_layers = self.get_fc_layers()
          self.n_fc = len(fc_layers)
          self.fc_model = to_gpu(nn.Sequential(*fc_layers))
          if not custom_head: apply_init(self.fc_model, kaiming_normal)
          self.model = to_gpu(nn.Sequential(*(layers + fc_layers)))
  
      @property
      def name(self):
          return f'{self.f.__name__}_{self.xtra_cut}'
  
      def create_fc_layer(self, ni, nf, p, actn=None):
          res = [nn.BatchNorm1d(num_features=ni)]
          if p: res.append(nn.Dropout(p=p))
          res.append(nn.Linear(in_features=ni, out_features=nf))
          if actn: res.append(actn)
          return res
  
      def get_fc_layers(self):
          res = []
          ni = self.nf
          for i, nf in enumerate(self.xtra_fc):
              res += self.create_fc_layer(ni, nf, p=self.ps[i], actn=nn.ReLU())
              ni = nf
          final_actn = nn.Sigmoid() if self.is_multi else nn.LogSoftmax()
          if self.is_reg: final_actn = None
          res += self.create_fc_layer(ni, self.c, p=self.ps[-1], actn=final_actn)
          return res
  
      def get_layer_groups(self, do_fc=False):
          if do_fc:
              return [self.fc_model]
          idxs = [self.lr_cut]
          c = children(self.top_model)
          if len(c) == 3: c = children(c[0]) + c[1:]
          lgs = list(split_by_idxs(c, idxs))
          return lgs + [self.fc_model]
  
  
  class ConvLearner(Learner):
      def __init__(self, data, models, precompute=False, **kwargs):
          self.precompute = False
          super().__init__(data, models, **kwargs)
          if hasattr(data, 'is_multi') and not data.is_reg and self.metrics is None:
              self.metrics = [accuracy_thresh(0.5)] if self.data.is_multi else [accuracy]
          if precompute: self.save_fc1()
          self.freeze()
          self.precompute = precompute
  
      def _get_crit(self, data):
          if not hasattr(data, 'is_multi'): return super()._get_crit(data)
  
          return F.l1_loss if data.is_reg else F.binary_cross_entropy if data.is_multi else F.nll_loss
  
      @classmethod
      def pretrained(cls, f, data, ps=None, xtra_fc=None, xtra_cut=0, custom_head=None, precompute=False,
                     pretrained=True, **kwargs):
          models = ConvnetBuilder_custom(f, data.c, data.is_multi, data.is_reg,
                                         ps=ps, xtra_fc=xtra_fc, xtra_cut=xtra_cut, custom_head=custom_head,
                                         pretrained=pretrained)
          return cls(data, models, precompute, **kwargs)
  
      @classmethod
      def lsuv_learner(cls, f, data, ps=None, xtra_fc=None, xtra_cut=0, custom_head=None, precompute=False,
                       needed_std=1.0, std_tol=0.1, max_attempts=10, do_orthonorm=False, **kwargs):
          models = ConvnetBuilder(f, data.c, data.is_multi, data.is_reg,
                                  ps=ps, xtra_fc=xtra_fc, xtra_cut=xtra_cut, custom_head=custom_head, pretrained=False)
          convlearn = cls(data, models, precompute, **kwargs)
          convlearn.lsuv_init()
          return convlearn
  
      @property
      def model(self):
          return self.models.fc_model if self.precompute else self.models.model
  
      def half(self):
          if self.fp16: return
          self.fp16 = True
          if type(self.model) != FP16: self.models.model = FP16(self.model)
          if not isinstance(self.models.fc_model, FP16): self.models.fc_model = FP16(self.models.fc_model)
  
      def float(self):
          if not self.fp16: return
          self.fp16 = False
          if type(self.models.model) == FP16: self.models.model = self.model.module.float()
          if type(self.models.fc_model) == FP16: self.models.fc_model = self.models.fc_model.module.float()
  
      @property
      def data(self):
          return self.fc_data if self.precompute else self.data_
  
      def create_empty_bcolz(self, n, name):
          return bcolz.carray(np.zeros((0, n), np.float32), chunklen=1, mode='w', rootdir=name)
  
      def set_data(self, data, precompute=False):
          super().set_data(data)
          if precompute:
              self.unfreeze()
              self.save_fc1()
              self.freeze()
              self.precompute = True
          else:
              self.freeze()
  
      def get_layer_groups(self):
          return self.models.get_layer_groups(self.precompute)
  
      def summary(self):
          precompute = self.precompute
          self.precompute = False
          res = super().summary()
          self.precompute = precompute
          return res
  
      def get_activations(self, force=False):
          tmpl = f'_{self.models.name}_{self.data.sz}.bc'
          # TODO: Somehow check that directory names haven't changed (e.g. added test set)
          names = [os.path.join(self.tmp_path, p + tmpl) for p in ('x_act', 'x_act_val', 'x_act_test')]
          if os.path.exists(names[0]) and not force:
              self.activations = [bcolz.open(p) for p in names]
          else:
              self.activations = [self.create_empty_bcolz(self.models.nf, n) for n in names]
  
      def save_fc1(self):
          self.get_activations()
          act, val_act, test_act = self.activations
          m = self.models.top_model
          if len(self.activations[0]) != len(self.data.trn_ds):
              predict_to_bcolz(m, self.data.fix_dl, act)
          if len(self.activations[1]) != len(self.data.val_ds):
              predict_to_bcolz(m, self.data.val_dl, val_act)
          if self.data.test_dl and (len(self.activations[2]) != len(self.data.test_ds)):
              if self.data.test_dl: predict_to_bcolz(m, self.data.test_dl, test_act)
  
          self.fc_data = ImageClassifierData.from_arrays(self.data.path,
                                                         (act, self.data.trn_y), (val_act, self.data.val_y), self.data.bs,
                                                         classes=self.data.classes,
                                                         test=test_act if self.data.test_dl else None, num_workers=8)
  
      def freeze(self):
          self.freeze_to(-1)
  
      def unfreeze(self):
          self.freeze_to(0)
          self.precompute = False
  
      def predict_array(self, arr):
          precompute = self.precompute
          self.precompute = False
          pred = super().predict_array(arr)
          self.precompute = precompute
          return pred
  
  
  preds_dict = {}
  preds_t_dict = {}
  
  # Training
  for fold_id in range(10):
      sz = 256 #image size
      bs = 64  #batch size
      tr_n, val_n = get_fold(fold_id,folds)
      md = get_data(sz,bs,tr_n, val_n)
      learner = ConvLearner.pretrained(arch, md, ps=0.5) #dropout 50%
      learner.models_path = MODEL_PATH
      learner.opt_fn = optim.Adam
      learner.clip = 1.0 #gradient clipping
      learner.crit = FocalLoss()
      learner.metrics = [acc]
  
      print(learner.summary)
  
      # find lr
      learner.lr_find()
      learner.sched.plot()
  
  
      # fit 1 epoch
      lr = 2e-2
      learner.fit(lr,1)
  
      #unfreeze and set differential lr
      learner.unfreeze()
      lrs=np.array([lr/10,lr/3,lr])
  
      learner.fit(lrs/4,4,cycle_len=2,use_clr=(10,20))
  
      learner.fit(lrs/4,2,cycle_len=4,use_clr=(10,20))
  
      learner.sched.plot_lr()
  
      learner.fit(lrs/16,1,cycle_len=8,use_clr=(5,20))
  
      # save model
      learner.save(f'ResNet34_256_{fold_id}')
  
      # evaluate on val set using tta
      def sigmoid_np(x):
          return 1.0/(1.0 + np.exp(-x))
  
      preds,y = learner.TTA(n_aug=16)
      preds = np.stack(preds, axis=-1)
      preds = sigmoid_np(preds)
      pred = preds.max(axis=-1)
  
      # optimize threshold
  
      def F1_soft(preds,targs,th=0.5,d=50.0):
          preds = sigmoid_np(d*(preds - th))
          targs = targs.astype(np.float)
          score = 2.0*(preds*targs).sum(axis=0)/((preds+targs).sum(axis=0) + 1e-6)
          return score
  
      def fit_val(x,y):
          params = 0.5*np.ones(len(name_label_dict))
          wd = 1e-5
          error = lambda p: np.concatenate((F1_soft(x,y,p) - 1.0,
                                            wd*(p - 0.5)), axis=None)
          p, success = opt.leastsq(error, params)
          return p
  
      th = fit_val(pred,y)
      th[th<0.1] = 0.1
      print('Thresholds: ',th)
      print('F1 macro: ',f1_score(y, pred>th, average='macro'))
      print('F1 macro (th = 0.5): ',f1_score(y, pred>0.5, average='macro'))
      print('F1 micro: ',f1_score(y, pred>th, average='micro'))
  
      print('Fractions: ',(pred > th).mean(axis=0))
      print('Fractions (true): ',(y > th).mean(axis=0))
  
  
      #populate submission
      preds_t,y_t = learner.TTA(n_aug=16,is_test=True)
      preds_t = np.stack(preds_t, axis=-1)
      preds_t = sigmoid_np(preds_t)
      pred_t = preds_t.max(axis=-1) #max works better for F1 macro score
  
      preds_dict[str(fold_id)] = pred
      preds_t_dict[str(fold_id)] = pred_t
  
  import pickle
  
  with open('preds_dict','wb') as f:
      pickle.dump(preds_dict,f)
  
  with open('preds_t_dict','wb') as f:
      pickle.dump(preds_t_dict,f)
  
  def save_pred(pred, th=0.5, fname='protein_classification.csv'):
      pred_list = []
      for line in pred:
          s = ' '.join(list([str(i) for i in np.nonzero(line > th)[0]]))
          pred_list.append(s)
  
      sample_df = pd.read_csv(SAMPLE)
      sample_list = list(sample_df.Id)
      pred_dic = dict((key, value) for (key, value)
                      in zip(learner.data.test_ds.fnames, pred_list))
      pred_list_cor = [pred_dic[id] for id in sample_list]
      df = pd.DataFrame({'Id': sample_list, 'Predicted': pred_list_cor})
      df.to_csv(fname, header=True, index=False)
  
  # manually found
  th_t = np.array([0.565,0.39,0.55,0.345,0.33,0.39,0.33,0.45,0.38,0.39,
                 0.34,0.42,0.31,0.38,0.49,0.50,0.38,0.43,0.46,0.40,
                 0.39,0.505,0.37,0.47,0.41,0.545,0.32,0.1])
  print('Fractions: ',(pred_t > th_t).mean(axis=0))
  save_pred(pred_t,th_t)
  
  
  # I replaced 0 by 0.01 since there may be a rounding error leading to 0
  
  def Count_soft(preds,th=0.5,d=50.0):
      preds = sigmoid_np(d*(preds - th))
      return preds.mean(axis=0)
  
  def fit_test(x,y):
      params = 0.5*np.ones(len(name_label_dict))
      wd = 1e-5
      error = lambda p: np.concatenate((Count_soft(x,p) - y,
                                        wd*(p - 0.5)), axis=None)
      p, success = opt.leastsq(error, params)
      return p
  
  pred_t = np.mean(np.array([preds_t_dict[str(f)] for f in range(10)]),axis=0)
  th_t = fit_test(pred_t,lb_prob)
  th_t[th_t<0.1] = 0.1
  print('Thresholds: ',th_t)
  print('Fractions: ',(pred_t > th_t).mean(axis=0))
  print('Fractions (th = 0.5): ',(pred_t > 0.5).mean(axis=0))
  
  save_pred(pred_t,th_t,'protein_classification_f.csv')
  
  save_pred(pred_t,th,'protein_classification_v.csv')
  save_pred(pred_t,0.5,'protein_classification_05.csv')
  
  from scipy.stats.mstats import gmean
  pred_t_geo = gmean(np.array([preds_t_dict[str(f)] for f in range(10)]),axis=0)
  th_t = fit_test(pred_t_geo,lb_prob)
  th_t[th_t<0.1] = 0.1
  print('Thresholds: ',th_t)
  print('Fractions: ',(pred_t_geo > th_t).mean(axis=0))
  print('Fractions (th = 0.5): ',(pred_t_geo > 0.5).mean(axis=0))
  
  save_pred(pred_t,th_t,'protein_classification_f_geo.csv')
  
  save_pred(pred_t,th,'protein_classification_v_geo.csv')
  save_pred(pred_t,0.5,'protein_classification_05_geo.csv')
  
  pred_t_max = np.max(np.array([preds_t_dict[str(f)] for f in range(10)]),axis=0)
  th_t = fit_test(pred_t_max,lb_prob)
  th_t[th_t<0.1] = 0.1
  print('Thresholds: ',th_t)
  print('Fractions: ',(pred_t_max > th_t).mean(axis=0))
  print('Fractions (th = 0.5): ',(pred_t_max > 0.5).mean(axis=0))
  
  save_pred(pred_t,th_t,'protein_classification_f_max.csv')
  
  save_pred(pred_t,th,'protein_classification_v_max.csv')
  save_pred(pred_t,0.5,'protein_classification_05_max.csv')
  
  
  
  tr_n, val_n = get_fold(0,folds)
  
  train_df = pd.read_csv('Christof/assets/train.csv')
  y_f = [train_df['']]
  
  
  th = fit_val(pred, y)
  th[th < 0.1] = 0.1
  print('Thresholds: ', th)
  print('F1 macro: ', f1_score(y, pred > th, average='macro'))
  print('F1 macro (th = 0.5): ', f1_score(y, pred > 0.5, average='macro'))
  print('F1 micro: ', f1_score(y, pred > th, average='micro'))
  
  
  
  
  class_list = [8,9,10,15,20,24,27]
  for i in class_list:
      th_t[i] = th[i]
  save_pred(pred_t,th_t,'protein_classification_c.csv')
  
  import matplotlib.pyplot as plt
  fig,ax = plt.subplots(2,1,figsize=(8,12))
  ax[0].plot(list(range(8)),learner.sched.val_losses, label='Validation loss')
  #ax[0].plot(list(range(8)),learner.sched.losses, label='Training loss')
  ax[0].set_xlabel('Epoch')
  ax[0].set_ylabel('Loss')
  ax[0].legend(loc='upper right')
  ax[1].plot(list(range(8)),learner.sched.rec_metrics)
  ax[1].set_xlabel('Epoch')
  ax[1].set_ylabel('Accuracy')