从图像标注开始，用 Byzer 加 Xtreme1 完成图像实例分割训练

#%python
#%input=new_ballon_dataset
#%output=out
#%schema=file
#%runIn=driver
#%dataMode=model
#%cache=true
#%env=source /opt/miniconda3/bin/activate ray-1.12.0


from pyjava.api.mlsql import RayContext,PythonContext
from tech.mlsql.plugin.ds.coco import create_coco_format_json

import torch, detectron2
import os

## Some basic setup:
## Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()


## import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data.datasets import register_coco_instances
from detectron2.engine import DefaultTrainer


context:PythonContext = context
ray_context = RayContext.connect(globals(),None)

rows = ray_context.collect()

create_coco_format_json(rows,"/tmp/new_ballon_dataset2")

TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
CUDA_VERSION = torch.__version__.split("+")[-1]
print("torch: ", TORCH_VERSION, "; cuda: ", CUDA_VERSION)
print("detectron2:", detectron2.__version__)


cfg = get_cfg()
## add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
## Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)

register_coco_instances("new_ballon_dataset", {}, "/tmp/new_ballon_dataset2/annotations/coco.json", "/tmp/new_ballon_dataset2/images")


cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("new_ballon_dataset",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2  ## This is the real "batch size" commonly known to deep learning people
cfg.SOLVER.BASE_LR = 0.00025  ## pick a good LR
cfg.SOLVER.MAX_ITER = 300    ## 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.SOLVER.STEPS = []        ## do not decay learning rate
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   ## The "RoIHead batch size". 128 is faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  ## only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
## NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

model_path = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") 
model_binary = [item for item in streaming_tar.build_rows_from_file(model_path)]

context.build_result(model_binary)