Spaces:
Runtime error
Runtime error
| import sys | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from PyQt5.QtCore import * | |
| from PyQt5.QtGui import * | |
| from PyQt5.QtWidgets import * | |
| from models.sample_model import SampleFromPoseModel | |
| from ui.mouse_event import GraphicsScene | |
| from ui.ui import Ui_Form | |
| from utils.language_utils import (generate_shape_attributes, | |
| generate_texture_attributes) | |
| from utils.options import dict_to_nonedict, parse | |
| color_list = [(0, 0, 0), (255, 250, 250), (220, 220, 220), (250, 235, 215), | |
| (255, 250, 205), (211, 211, 211), (70, 130, 180), | |
| (127, 255, 212), (0, 100, 0), (50, 205, 50), (255, 255, 0), | |
| (245, 222, 179), (255, 140, 0), (255, 0, 0), (16, 78, 139), | |
| (144, 238, 144), (50, 205, 174), (50, 155, 250), (160, 140, 88), | |
| (213, 140, 88), (90, 140, 90), (185, 210, 205), (130, 165, 180), | |
| (225, 141, 151)] | |
| class Ex(QWidget, Ui_Form): | |
| def __init__(self, opt): | |
| super(Ex, self).__init__() | |
| self.setupUi(self) | |
| self.show() | |
| self.output_img = None | |
| self.mat_img = None | |
| self.mode = 0 | |
| self.size = 6 | |
| self.mask = None | |
| self.mask_m = None | |
| self.img = None | |
| # about UI | |
| self.mouse_clicked = False | |
| self.scene = QGraphicsScene() | |
| self.graphicsView.setScene(self.scene) | |
| self.graphicsView.setAlignment(Qt.AlignTop | Qt.AlignLeft) | |
| self.graphicsView.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) | |
| self.graphicsView.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) | |
| self.ref_scene = GraphicsScene(self.mode, self.size) | |
| self.graphicsView_2.setScene(self.ref_scene) | |
| self.graphicsView_2.setAlignment(Qt.AlignTop | Qt.AlignLeft) | |
| self.graphicsView_2.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) | |
| self.graphicsView_2.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) | |
| self.result_scene = QGraphicsScene() | |
| self.graphicsView_3.setScene(self.result_scene) | |
| self.graphicsView_3.setAlignment(Qt.AlignTop | Qt.AlignLeft) | |
| self.graphicsView_3.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) | |
| self.graphicsView_3.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) | |
| self.dlg = QColorDialog(self.graphicsView) | |
| self.color = None | |
| self.sample_model = SampleFromPoseModel(opt) | |
| def open_densepose(self): | |
| fileName, _ = QFileDialog.getOpenFileName(self, "Open File", | |
| QDir.currentPath()) | |
| if fileName: | |
| image = QPixmap(fileName) | |
| mat_img = Image.open(fileName) | |
| self.pose_img = mat_img.copy() | |
| if image.isNull(): | |
| QMessageBox.information(self, "Image Viewer", | |
| "Cannot load %s." % fileName) | |
| return | |
| image = image.scaled(self.graphicsView.size(), | |
| Qt.IgnoreAspectRatio) | |
| if len(self.scene.items()) > 0: | |
| self.scene.removeItem(self.scene.items()[-1]) | |
| self.scene.addPixmap(image) | |
| self.ref_scene.clear() | |
| self.result_scene.clear() | |
| # load pose to model | |
| self.pose_img = np.array( | |
| self.pose_img.resize( | |
| size=(256, 512), | |
| resample=Image.LANCZOS))[:, :, 2:].transpose( | |
| 2, 0, 1).astype(np.float32) | |
| self.pose_img = self.pose_img / 12. - 1 | |
| self.pose_img = torch.from_numpy(self.pose_img).unsqueeze(1) | |
| self.sample_model.feed_pose_data(self.pose_img) | |
| def generate_parsing(self): | |
| self.ref_scene.reset_items() | |
| self.ref_scene.reset() | |
| shape_texts = self.message_box_1.text() | |
| shape_attributes = generate_shape_attributes(shape_texts) | |
| shape_attributes = torch.LongTensor(shape_attributes).unsqueeze(0) | |
| self.sample_model.feed_shape_attributes(shape_attributes) | |
| self.sample_model.generate_parsing_map() | |
| self.sample_model.generate_quantized_segm() | |
| self.colored_segm = self.sample_model.palette_result( | |
| self.sample_model.segm[0].cpu()) | |
| self.mask_m = cv2.cvtColor( | |
| cv2.cvtColor(self.colored_segm, cv2.COLOR_RGB2BGR), | |
| cv2.COLOR_BGR2RGB) | |
| qim = QImage(self.colored_segm.data.tobytes(), | |
| self.colored_segm.shape[1], self.colored_segm.shape[0], | |
| QImage.Format_RGB888) | |
| image = QPixmap.fromImage(qim) | |
| image = image.scaled(self.graphicsView.size(), Qt.IgnoreAspectRatio) | |
| if len(self.ref_scene.items()) > 0: | |
| self.ref_scene.removeItem(self.ref_scene.items()[-1]) | |
| self.ref_scene.addPixmap(image) | |
| self.result_scene.clear() | |
| def generate_human(self): | |
| for i in range(24): | |
| self.mask_m = self.make_mask(self.mask_m, | |
| self.ref_scene.mask_points[i], | |
| self.ref_scene.size_points[i], | |
| color_list[i]) | |
| seg_map = np.full(self.mask_m.shape[:-1], -1) | |
| # convert rgb to num | |
| for index, color in enumerate(color_list): | |
| seg_map[np.sum(self.mask_m == color, axis=2) == 3] = index | |
| assert (seg_map != -1).all() | |
| self.sample_model.segm = torch.from_numpy(seg_map).unsqueeze( | |
| 0).unsqueeze(0).to(self.sample_model.device) | |
| self.sample_model.generate_quantized_segm() | |
| texture_texts = self.message_box_2.text() | |
| texture_attributes = generate_texture_attributes(texture_texts) | |
| texture_attributes = torch.LongTensor(texture_attributes) | |
| self.sample_model.feed_texture_attributes(texture_attributes) | |
| self.sample_model.generate_texture_map() | |
| result = self.sample_model.sample_and_refine() | |
| result = result.permute(0, 2, 3, 1) | |
| result = result.detach().cpu().numpy() | |
| result = result * 255 | |
| result = np.asarray(result[0, :, :, :], dtype=np.uint8) | |
| self.output_img = result | |
| qim = QImage(result.data.tobytes(), result.shape[1], result.shape[0], | |
| QImage.Format_RGB888) | |
| image = QPixmap.fromImage(qim) | |
| image = image.scaled(self.graphicsView.size(), Qt.IgnoreAspectRatio) | |
| if len(self.result_scene.items()) > 0: | |
| self.result_scene.removeItem(self.result_scene.items()[-1]) | |
| self.result_scene.addPixmap(image) | |
| def top_mode(self): | |
| self.ref_scene.mode = 1 | |
| def skin_mode(self): | |
| self.ref_scene.mode = 15 | |
| def outer_mode(self): | |
| self.ref_scene.mode = 2 | |
| def face_mode(self): | |
| self.ref_scene.mode = 14 | |
| def skirt_mode(self): | |
| self.ref_scene.mode = 3 | |
| def hair_mode(self): | |
| self.ref_scene.mode = 13 | |
| def dress_mode(self): | |
| self.ref_scene.mode = 4 | |
| def headwear_mode(self): | |
| self.ref_scene.mode = 7 | |
| def pants_mode(self): | |
| self.ref_scene.mode = 5 | |
| def eyeglass_mode(self): | |
| self.ref_scene.mode = 8 | |
| def rompers_mode(self): | |
| self.ref_scene.mode = 21 | |
| def footwear_mode(self): | |
| self.ref_scene.mode = 11 | |
| def leggings_mode(self): | |
| self.ref_scene.mode = 6 | |
| def ring_mode(self): | |
| self.ref_scene.mode = 16 | |
| def belt_mode(self): | |
| self.ref_scene.mode = 10 | |
| def neckwear_mode(self): | |
| self.ref_scene.mode = 9 | |
| def wrist_mode(self): | |
| self.ref_scene.mode = 17 | |
| def socks_mode(self): | |
| self.ref_scene.mode = 18 | |
| def tie_mode(self): | |
| self.ref_scene.mode = 23 | |
| def earstuds_mode(self): | |
| self.ref_scene.mode = 22 | |
| def necklace_mode(self): | |
| self.ref_scene.mode = 20 | |
| def bag_mode(self): | |
| self.ref_scene.mode = 12 | |
| def glove_mode(self): | |
| self.ref_scene.mode = 19 | |
| def background_mode(self): | |
| self.ref_scene.mode = 0 | |
| def make_mask(self, mask, pts, sizes, color): | |
| if len(pts) > 0: | |
| for idx, pt in enumerate(pts): | |
| cv2.line(mask, pt['prev'], pt['curr'], color, sizes[idx]) | |
| return mask | |
| def save_img(self): | |
| if type(self.output_img): | |
| fileName, _ = QFileDialog.getSaveFileName(self, "Save File", | |
| QDir.currentPath()) | |
| cv2.imwrite(fileName + '.png', self.output_img[:, :, ::-1]) | |
| def undo(self): | |
| self.scene.undo() | |
| def clear(self): | |
| self.ref_scene.reset_items() | |
| self.ref_scene.reset() | |
| self.ref_scene.clear() | |
| self.result_scene.clear() | |
| if __name__ == '__main__': | |
| app = QApplication(sys.argv) | |
| opt = './configs/sample_from_pose.yml' | |
| opt = parse(opt, is_train=False) | |
| opt = dict_to_nonedict(opt) | |
| ex = Ex(opt) | |
| sys.exit(app.exec_()) | |