From 70a42c21b573f2de78a7007f8d72855d10a0e9e1 Mon Sep 17 00:00:00 2001
From: Hunor Laczko <luni741@gmail.com>
Date: Thu, 22 Dec 2022 23:16:45 +0000
Subject: [PATCH] generate registered, and frames_file

---
 .gitignore         |  3 ++-
 data/preprocess.py | 63 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index d4be5a2..9969755 100644
--- a/.gitignore
+++ b/.gitignore
@@ -161,4 +161,5 @@ cython_debug/
 #.idea/
 
 .vscode/
-temp/
\ No newline at end of file
+temp/
+dataset/
\ No newline at end of file
diff --git a/data/preprocess.py b/data/preprocess.py
index 2b6c934..6a856d1 100644
--- a/data/preprocess.py
+++ b/data/preprocess.py
@@ -2,10 +2,13 @@ import os
 import sys
 
 import numpy as np
+import scipy as sp
+import scipy.io as sio
 from tqdm import tqdm
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from utils.utils import mask_faces
+from superSMPL.IO import readPC2, writePC2
 
 def generate_faces(folder):
     frames = os.listdir(os.path.join(folder, 'mesh_mask'))
@@ -18,8 +21,64 @@ def generate_faces(folder):
         np.save(os.path.join(folder, 'faces', frame), faces)
 
 
+class Frame:
+    def __init__(self, id, frame_nr):
+        self.id = id
+        self.frame_nr = frame_nr
+
+
+def generate_frames_file(folder, file_name):
+    ids = os.listdir(folder)#[:1]
+    frames = []
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+        garments = [f.replace('_registered.pc16','') for f in os.listdir(garment_dir) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
+        assert len(garments) == 1
+        garment = garments[0]
+        verts = readPC2(os.path.join(garment_dir, garment + '_registered.pc16'), True)['V']
+        for i in range(verts.shape[0]):
+            frames.append(Frame(id, i))
+
+    with open(os.path.join(folder, file_name), 'w') as f:
+        for i, frame in enumerate(frames):
+            f.write(str(frame.id) + ' ' + str(frame.frame_nr) + '\n')
+
+
+def genereated_registered_and_mask(folder, is_skirt = False):
+    ids = os.listdir(folder)
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')]
+        
+        assert len(garments) == 1
+
+        if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')):
+            continue
+        
+        garment = garments[0]
+
+        verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
+        map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1))
+        verts = verts[:, map, :] # ---------------> This is the garment after registration
+        raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1))
+        mask = raw_mask > 0
+        # boundaries = raw_mask == 2
+
+        writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True)
+        np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask)
+
+
 if __name__ == '__main__':
     # folder = '/home/workspace/local-data/train'
-    folder = '/home/workspace/local-data/valid'
+    # folder = '/home/workspace/local-data/valid'
+    folder = 'dataset/CLOTH3D/train'
 
-    generate_faces(folder)
+    # generate_faces(folder)
+    # genereated_registered_and_mask(folder)
+    generate_frames_file(folder, 'frames_all.txt')