/
MixAndMatchApp.py
137 lines (113 loc) · 4.32 KB
/
MixAndMatchApp.py
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
import progressbar as pbar #pip3 install progressbar2
import numpy as np
import os
import random
import sys
import trimesh
import DatasetUtils as dt
import DatasetViewerApp as dva
import ModelPartsScreenshot as mps
import ModelPartsViewer as mpv
import ProjectUtils as pUtils
from Scorer import Scorer
from MLStatics import *
if __name__ == "__main__":
tokens = sys.argv[1:]
if len(tokens) <= 0:
print('- Run with random dataset indices e.g. \"MixAndMatchApp.py -r 5\"')
print('- Run with defense sets e.g. \"MixAndMatchApp.py -setA (or setB, setC)\"')
print('- Specify how many chairs to generate and the pool size e.g. \"MixAndMatchApp.py -r 5 -g 10 -p 100\"')
quit()
chairsToGenerateCount = 10
if '-g' in tokens:
index = tokens.index('-g')
chairsToGenerateCount = int(tokens[index+1])
chairPoolCount = 200
if '-p' in tokens:
index = tokens.index('-p')
chairPoolCount = int(tokens[index+1])
datasetIndices = []
if '-r' in tokens:
index = tokens.index('-r')
randomAmount = int(tokens[index+1])
randomShuffledIndices = list(range(1, 6202))
random.shuffle(randomShuffledIndices)
datasetIndices = randomShuffledIndices[:randomAmount]
if '-setA' in tokens:
datasetIndices = ['369', '175', '5540']
if '-setB' in tokens:
datasetIndices = ['2999', '2150', '3492', '4474', '2160']
if '-setC' in tokens:
datasetIndices = ['1919', '3366', '3521', '3204',
'1131', '173', '3749', '2313', '5117', '1920']
# conver to str
inputIndices = []
for index in datasetIndices:
inputIndices.append(str(index))
# Start
print('Running program for indices: '+str(inputIndices))
print('# of chairs to generate: '+str(chairsToGenerateCount))
print('Size of the pool: '+str(chairPoolCount))
models = []
for index in inputIndices:
partsTuples = dt.getDatasetObjParts(index)
modelParts = dva.parseDatasetChairTuples(index, partsTuples)
model = mpv.Model(modelParts)
model.name = str(index)
model.datasetIndex = str(index)
model.datasetObjIndex = dt.getDatasetObjIndex(index)
models.append(model)
# fill the chair pool, do not repeat the chairs
newChairs = []
print("Generating Chairs...")
tryGenerateTimes = 100
tryGenerateCount = 0
bar = pbar.ProgressBar(max_value = chairPoolCount)
while len(newChairs) < chairPoolCount:
newChair = mpv.generateChairFast(models)
if mpv.modelListContains(newChairs, newChair):
tryGenerateCount+=1
if tryGenerateCount >= tryGenerateTimes:
break
continue
mpv.weldChairParts(newChair)
tryGenerateCount = 0
newChairs.append(newChair)
bar.update(len(newChairs))
bar.finish()
# screenshot every new chair
rotations = [
(0, 0, 0), # front
(0, np.pi/2, 0), # right
(np.pi/2, 0, 0), # top
# (0, np.pi, 0), # back
# (0, -np.pi/2, 0), # left
# (-np.pi/2, 0, 0) # bottom
]
depthScreenshots = []
for generatedChair in progressbar(newChairs, "Making Screenshots"):
perspectives = mps.captureDepth(
generatedChair, rotations, imageWidth=224, imageHeight=224)
depthScreenshots.append((generatedChair, perspectives))
s = Scorer()
# Assign a score
scoredChairs = []
for generatedChair, perspectives in progressbar(depthScreenshots, "Evaluating Chairs"):
score = s.score(perspectives)
generatedChair.cachedScore = score
scoredChairs.append((generatedChair, score))
# sort models depending on the score, from bigger to smaller
chairsToDisplay = []
sortedChairs = sorted(scoredChairs, reverse=True, key=lambda tup: tup[1])
for chair, value in sortedChairs:
print(value)
chairsToDisplay.append(chair)
# we would need only 10 top chairs
chairsToDisplay = chairsToDisplay[:chairsToGenerateCount]
# export the chairs as .objs
for i in progressbar(range(len(chairsToDisplay)), "Exporting generated chairs as .obj"):
fileName = str(i) + '.obj'
mpv.exportModelAsObj(chairsToDisplay[i], fileName)
# set models in a viewer
mpv.setInputModels(chairsToDisplay)
mpv.start()