Skip to content

Commit

Permalink
Update detector.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@yujiahu415
yujiahu415 committed Nov 7, 2024
1 parent 81179c6 commit 7ba532b
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions LabGym/detector.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,12 +79,12 @@ def train(self,path_to_annotation,path_to_trainingimages,path_to_detector,iterat
print('Animal names in annotation file: '+str(model_parameters_dict['animal_names']))

cfg=get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.merge_from_file(model_zoo.get_config_file('COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml'))
cfg.OUTPUT_DIR=path_to_detector
cfg.DATASETS.TRAIN=('LabGym_detector_train',)
cfg.DATASETS.TEST=()
cfg.DATALOADER.NUM_WORKERS=4
cfg.MODEL.WEIGHTS=model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS=model_zoo.get_checkpoint_url('COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE=128
cfg.MODEL.ROI_HEADS.NUM_CLASSES=int(len(classnames))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST=0.5
Expand Down

0 comments on commit 7ba532b

Please sign in to comment.