-
Notifications
You must be signed in to change notification settings - Fork 0
/
ObjectLabler_v1.cs
154 lines (127 loc) · 5.48 KB
/
ObjectLabler_v1.cs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
using System.Collections;
using System.Collections.Generic;
using System.IO;
using UnityEngine;
public class ObjectLabeler : MonoBehaviour
{
public Camera segmentationCamera; // Segmentation 마스크를 위한 별도의 카메라
public string outputDirectory = "Assets/labeledScreenshot/";
public int imageWidth = 1024;
public int imageHeight = 1024;
public float preRenderDelay = 0.1f; // 색상 변경 후 대기 시간
public float postRenderDelay = 0.1f; // 스크린샷 후 대기 시간
public List<GameObject> objectsToLabel = new List<GameObject>(); // 라벨링할 객체들을 담는 리스트
private Dictionary<string, Color> objectColorMapping = new Dictionary<string, Color>
{
{ "longi", Color.red },
{ "floor", Color.green },
{ "realFloor", Color.blue },
{ "plate", Color.yellow },
{ "slotHole", Color.magenta },
{ "r_hole", Color.cyan }
};
public LightController lightController; // LightController를 참조하는 변수 추가
public void LabelObjects(string screenshotFilePath)
{
// 원본 이미지 촬영 후 텍스처와 조명을 제거한 상태에서 마스크 이미지를 촬영하도록 설정
string directory = outputDirectory;
string fileName = Path.GetFileNameWithoutExtension(screenshotFilePath) + "_mask.png";
string maskFilePath = Path.Combine(directory, fileName);
// 텍스처 제거 및 조명 재설정
RemoveTexturesAndAdjustLighting();
// slotHole과 r_hole의 위치 조정
AdjustObjectPositions();
// 마스크 생성
StartCoroutine(RenderFullMask(maskFilePath));
}
void GenerateSegmentationMasksAndLabels()
{
foreach (KeyValuePair<string, Color> entry in objectColorMapping)
{
string tag = entry.Key;
Color color = entry.Value;
GameObject[] objects = GameObject.FindGameObjectsWithTag(tag); // 태그로 객체를 식별
if (objects.Length == 0)
{
Debug.LogWarning($"No objects found with tag: {tag}");
continue;
}
SetObjectColors(objects, color);
// 라벨링할 객체 리스트에 추가
objectsToLabel.AddRange(objects);
}
}
void SetObjectColors(GameObject[] objects, Color color)
{
foreach (GameObject obj in objects)
{
if (!obj.activeInHierarchy)
{
obj.SetActive(true); // 오브젝트 활성화
}
Renderer[] renderers = obj.GetComponentsInChildren<Renderer>();
foreach (Renderer renderer in renderers)
{
renderer.material = new Material(Shader.Find("Unlit/Color")); // 새로운 단일 색상 머티리얼 설정
renderer.material.color = color; // 색상 설정
}
}
}
void RemoveTexturesAndAdjustLighting()
{
// 조명 설정을 변경하여 그림자를 제거하고 밝기를 높임
if (lightController != null)
{
lightController.light1.DisableShadows();
lightController.light1.SetIntensity(2.5f); // 세그멘테이션 마스크를 위한 밝기 조정
}
}
public void AdjustObjectPositions()
{
// slotHole과 r_hole의 위치를 약간 앞으로 이동
for (int i = objectsToLabel.Count - 1; i >= 0; i--)
{
GameObject obj = objectsToLabel[i];
if (obj == null)
{
// null 참조를 방지하기 위해 null인 오브젝트는 리스트에서 제거
objectsToLabel.RemoveAt(i);
continue;
}
if (obj.CompareTag("slotHole"))
{
obj.transform.position += new Vector3(0.05f, 0, 0); // x축으로 0.05만큼 이동
}
if (obj.CompareTag("r_hole"))
{
obj.transform.position += new Vector3(0.08f, 0, 0); // x축으로 0.08만큼 이동
}
}
}
IEnumerator RenderFullMask(string maskFilePath)
{
// 색상 변경 및 조명 조정 후 일정 시간 대기
GenerateSegmentationMasksAndLabels();
yield return new WaitForSeconds(preRenderDelay);
// RenderTexture 생성 및 카메라에 할당
RenderTexture renderTexture = new RenderTexture(imageWidth, imageHeight, 24);
segmentationCamera.targetTexture = renderTexture;
segmentationCamera.Render();
// 카메라 렌더링이 완료되었는지 확인한 후 스크린샷 촬영
yield return new WaitForEndOfFrame();
// RenderTexture를 텍스처로 변환
RenderTexture.active = renderTexture;
Texture2D maskTexture = new Texture2D(imageWidth, imageHeight, TextureFormat.RGB24, false);
maskTexture.ReadPixels(new Rect(0, 0, imageWidth, imageHeight), 0, 0);
maskTexture.Apply();
RenderTexture.active = null;
// 마스크 이미지 저장 (스크린샷 파일 이름과 동일하게 저장)
byte[] bytes = maskTexture.EncodeToPNG();
File.WriteAllBytes(maskFilePath, bytes);
Debug.Log($"Saved segmentation mask to {maskFilePath}");
// Clean up
segmentationCamera.targetTexture = null;
renderTexture.Release(); // RenderTexture.Release를 사용하여 메모리에서 해제
yield return new WaitForSeconds(postRenderDelay); // 스크린샷 후 대기
}
}