| 402 | |
| 403 | |
| 404 | class BodySplit(object): |
| 405 | def __init__(self,extra_info=None,bg_type=0,aug_type=-1,split_prob=0.5): |
| 406 | """ |
| 407 | aug: for half augmentation |
| 408 | -1: no aug |
| 409 | 0:only return aug img wo padding |
| 410 | 1:add padding |
| 411 | """ |
| 412 | self.extra_info=extra_info |
| 413 | self.bg_type=bg_type |
| 414 | self.mean=[104,116,124] |
| 415 | self.aug_type=aug_type |
| 416 | self.debug=1 |
| 417 | self.split_prob=split_prob |
| 418 | if extra_info is not None: |
| 419 | print('read extra_info') |
| 420 | if not os.path.exists(extra_info): |
| 421 | extra_info = extra_info.replace('/mnt/lustre/share', '/mnt/lustre/share_data') # sh1986 |
| 422 | f = open(extra_info, 'r') |
| 423 | self.info = dict() |
| 424 | for line in f.readlines(): |
| 425 | items = line.strip('\n').split(' ') |
| 426 | key = items[0]#.split('/')[-1] |
| 427 | #print(key) |
| 428 | #import pdb;pdb.set_trace() |
| 429 | self.info[key] = (float(items[-4]), float(items[-3]), float(items[-2]), float(items[-1])) |
| 430 | f.close() |
| 431 | print('Done') |
| 432 | |
| 433 | def blind_split(self, raw_img): |
| 434 | w, h = raw_img.size |
| 435 | if np.random.rand() <= self.split_prob: |
| 436 | if h > w * 0.7: |
| 437 | mid = max(w*0.7, int(round(h*0.35))) |
| 438 | cut = np.random.randint(low=mid, high=h) |
| 439 | # raw_img = raw_img[:cut] |
| 440 | raw_img = raw_img.crop((0, 0, w, cut)) |
| 441 | return raw_img |
| 442 | |
| 443 | def __call__(self, raw_img,fname): |
| 444 | if self.extra_info is None: |
| 445 | return raw_img |
| 446 | #name = fname.split('/')[-1] |
| 447 | |
| 448 | if fname not in self.info.keys(): |
| 449 | print("="*80) |
| 450 | print(fname) |
| 451 | return raw_img |
| 452 | |
| 453 | head, _, mid, foot = self.info[fname] |
| 454 | if head<0: |
| 455 | return raw_img |
| 456 | if self.aug_type==0: #crop from ori img |
| 457 | if np.random.rand() <= self.split_prob and mid<=1 and mid>0.15: |
| 458 | h, w, c = raw_img.shape |
| 459 | mid=int(mid*h) |
| 460 | raw_img=raw_img[:mid] |
| 461 | if self.debug: |