jeonghyo.k
3946ff6a25
- prediction/image/ FastAPI 서버 Docker 환경 구성 - Dockerfile: PyTorch 2.1 + CUDA 12.1 기반 GPU 이미지 - docker-compose.yml: GPU 할당 + 데이터 볼륨 마운트 - requirements.txt: 서버 의존성 목록 - .env.example: 환경변수 템플릿 - DOCKER_USAGE.md: 빌드/실행/API 사용법 문서 - Dockerfile에 .dockerignore 제외 폴더 mkdir -p 추가 - .gitignore: prediction/image 결과물 및 모델 가중치(.pth) 제외 추가 - dbInsert_csv.py, dbInsert_shp.py 삭제 (미사용 DB 로직) - api.py: dbInsert import 및 주석 처리된 DB 호출 코드 제거 - aerialRouter.ts: req.params 타입 오류 수정
72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import argparse
|
|
import os.path as osp
|
|
|
|
import mmcv
|
|
import torch
|
|
from mmcv.runner import CheckpointLoader
|
|
|
|
|
|
def convert_stdc(ckpt, stdc_type):
|
|
new_state_dict = {}
|
|
if stdc_type == 'STDC1':
|
|
stage_lst = ['0', '1', '2.0', '2.1', '3.0', '3.1', '4.0', '4.1']
|
|
else:
|
|
stage_lst = [
|
|
'0', '1', '2.0', '2.1', '2.2', '2.3', '3.0', '3.1', '3.2', '3.3',
|
|
'3.4', '4.0', '4.1', '4.2'
|
|
]
|
|
for k, v in ckpt.items():
|
|
ori_k = k
|
|
flag = False
|
|
if 'cp.' in k:
|
|
k = k.replace('cp.', '')
|
|
if 'features.' in k:
|
|
num_layer = int(k.split('.')[1])
|
|
feature_key_lst = 'features.' + str(num_layer) + '.'
|
|
stages_key_lst = 'stages.' + stage_lst[num_layer] + '.'
|
|
k = k.replace(feature_key_lst, stages_key_lst)
|
|
flag = True
|
|
if 'conv_list' in k:
|
|
k = k.replace('conv_list', 'layers')
|
|
flag = True
|
|
if 'avd_layer.' in k:
|
|
if 'avd_layer.0' in k:
|
|
k = k.replace('avd_layer.0', 'downsample.conv')
|
|
elif 'avd_layer.1' in k:
|
|
k = k.replace('avd_layer.1', 'downsample.bn')
|
|
flag = True
|
|
if flag:
|
|
new_state_dict[k] = ckpt[ori_k]
|
|
|
|
return new_state_dict
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description='Convert keys in official pretrained STDC1/2 to '
|
|
'MMSegmentation style.')
|
|
parser.add_argument('src', help='src model path')
|
|
# The dst path must be a full path of the new checkpoint.
|
|
parser.add_argument('dst', help='save path')
|
|
parser.add_argument('type', help='model type: STDC1 or STDC2')
|
|
args = parser.parse_args()
|
|
|
|
checkpoint = CheckpointLoader.load_checkpoint(args.src, map_location='cpu')
|
|
if 'state_dict' in checkpoint:
|
|
state_dict = checkpoint['state_dict']
|
|
elif 'model' in checkpoint:
|
|
state_dict = checkpoint['model']
|
|
else:
|
|
state_dict = checkpoint
|
|
|
|
assert args.type in ['STDC1',
|
|
'STDC2'], 'STD type should be STDC1 or STDC2!'
|
|
weight = convert_stdc(state_dict, args.type)
|
|
mmcv.mkdir_or_exist(osp.dirname(args.dst))
|
|
torch.save(weight, args.dst)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|