Coverage for src/bioimageio/spec/model/_v0_3_converter.py: 6%

1# type: ignore 

2from __future__ import annotations 

3 

4from typing import Any 

5 

6 

7def convert_model_from_v0_3_to_0_4_0(data: dict[str, Any]) -> None: 

8 """auto converts model 'data' to newest format""" 

9 

10 if "format_version" not in data: 

11 return 

12 

13 if data["format_version"] == "0.3.0": 

14 # no breaking change, bump to 0.3.1 

15 data["format_version"] = "0.3.1" 

16 

17 if data["format_version"] == "0.3.1": 

18 data = _convert_model_v0_3_1_to_v0_3_2(data) 

19 

20 if data["format_version"] == "0.3.2": 

21 data = _convert_model_v0_3_2_to_v0_3_3(data) 

22 

23 if data["format_version"] in ("0.3.3", "0.3.4", "0.3.5"): 

24 data["format_version"] = "0.3.6" 

25 

26 if data["format_version"] != "0.3.6": 

27 return 

28 

29 # remove 'future' from config if no other than the used future entries exist 

30 config = data.get("config", {}) 

31 if config.get("future") == {}: 

32 del config["future"] 

33 

34 # remove 'config' if now empty 

35 if data.get("config") == {}: 

36 del data["config"] 

37 

38 data.pop("language", None) 

39 data.pop("framework", None) 

40 

41 architecture = data.pop("source", None) 

42 architecture_sha256 = data.pop("sha256", None) 

43 kwargs = data.pop("kwargs", None) 

44 pytorch_state_dict_weights_entry = data.get("weights", {}).get("pytorch_state_dict") 

45 if pytorch_state_dict_weights_entry is not None: 

46 if architecture is not None: 

47 pytorch_state_dict_weights_entry["architecture"] = architecture 

48 

49 if architecture_sha256 is not None: 

50 pytorch_state_dict_weights_entry["architecture_sha256"] = ( 

51 architecture_sha256 

52 ) 

53 

54 if kwargs is not None: 

55 pytorch_state_dict_weights_entry["kwargs"] = kwargs 

56 

57 torchscript_weights_entry = data.get("weights", {}).pop("pytorch_script", None) 

58 if torchscript_weights_entry is not None: 

59 data.setdefault("weights", {})["torchscript"] = torchscript_weights_entry 

60 

61 data["format_version"] = "0.4.0" 

62 

63 

64def _convert_model_v0_3_1_to_v0_3_2(data: dict[str, Any]) -> dict[str, Any]: 

65 data["type"] = "model" 

66 data["format_version"] = "0.3.2" 

67 future = data.get("config", {}).get("future", {}).pop("0.3.2", {}) 

68 

69 authors = data.get("authors") 

70 if isinstance(authors, list): 

71 data["authors"] = [{"name": name} for name in authors] 

72 authors_update = future.get("authors") 

73 if authors_update is not None: 

74 for a, u in zip(data["authors"], authors_update): 

75 a.update(u) 

76 

77 # packaged_by 

78 packaged_by = data.get("packaged_by") 

79 if packaged_by is not None: 

80 data["packaged_by"] = [{"name": name} for name in data["packaged_by"]] 

81 packaged_by_update = future.get("packaged_by") 

82 if packaged_by_update is not None: 

83 for a, u in zip(data["packaged_by"], packaged_by_update): 

84 a.update(u) 

85 

86 # authors of weights 

87 weights = data.get("weights") 

88 if isinstance(weights, dict): 

89 for weights_format, weights_entry in weights.items(): 

90 if "authors" not in weights_entry: 

91 continue 

92 

93 weights_entry["authors"] = [ 

94 {"name": name} for name in weights_entry["authors"] 

95 ] 

96 authors_update = ( 

97 future.get("weights", {}).get(weights_format, {}).get("authors") 

98 ) 

99 if authors_update is not None: 

100 for a, u in zip(weights_entry["authors"], authors_update): 

101 a.update(u) 

102 

103 # model version 

104 if "version" in future: 

105 data["version"] = future.pop("version") 

106 

107 return data 

108 

109 

110def _convert_model_v0_3_2_to_v0_3_3(data: dict[str, Any]) -> dict[str, Any]: 

111 data["format_version"] = "0.3.3" 

112 if "outputs" in data: 

113 for out in data["outputs"]: 

114 if "shape" in out: 

115 shape = out["shape"] 

116 if isinstance(shape, dict) and "reference_input" in shape: 

117 shape["reference_tensor"] = shape.pop("reference_input") 

118 

119 return data