class Ex(QtWidgets.QWidget, Ui_Form): def __init__(self, model, config): super().__init__() #self.get_head_outline() self._step_counter = 1 # start camera self.record_video = RecordVideo() # connect the frame data signal and slot together self.record_video.frame_data.connect(self.camera_data_slot) # start face detector self.detector = MTCNNDetector(device='cpu') self.setupUi(self) self.show() self.model = model self.config = config self.model.load_demo_graph(config) self.output_img = None self.mat_img = None self.ld_mask = None self.ld_sk = None self._frame_data = None self.modes = [0, 0, 0] self.mouse_clicked = False self.scene = GraphicsScene(self.modes) self.graphicsView.setScene(self.scene) self.graphicsView.setAlignment(QtCore.Qt.AlignTop | QtCore.Qt.AlignLeft) self.graphicsView.setVerticalScrollBarPolicy( QtCore.Qt.ScrollBarAlwaysOff) self.graphicsView.setHorizontalScrollBarPolicy( QtCore.Qt.ScrollBarAlwaysOff) self.result_scene = QtWidgets.QGraphicsScene() self.graphicsView_2.setScene(self.result_scene) self.graphicsView_2.setAlignment(QtCore.Qt.AlignTop | QtCore.Qt.AlignLeft) self.graphicsView_2.setVerticalScrollBarPolicy( QtCore.Qt.ScrollBarAlwaysOff) self.graphicsView_2.setHorizontalScrollBarPolicy( QtCore.Qt.ScrollBarAlwaysOff) self.dlg = QtWidgets.QColorDialog(self.graphicsView) self.color = None def get_head_outline(self): """load head outline.""" head_outline = cv2.imread('./ui/head_outline.jpg') head_outline = head_outline[25:-125, 75:-75, :] head_outline = cv2.resize(head_outline, (512, 512), interpolation=cv2.INTER_CUBIC) head_outline[head_outline < 75] = 0.5 head_outline[head_outline > 1] = 1 self._head_outline = head_outline def mode_select(self, mode): for i in range(len(self.modes)): self.modes[i] = 0 self.modes[mode] = 1 def camera_data_slot(self, frame_data): self._step_counter += 1 self._frame_data = frame_data # resize to 512x512 frame_data = cv2.resize(frame_data, (512, 512), interpolation=cv2.INTER_CUBIC) # face detection # BGR to RGB first im = Image.fromarray(frame_data[:, :, ::-1]) bboxes, landmarks = self.detector.infer(im, min_face_size=200) if len(bboxes): _img = show_grids(im, [], landmarks, step=self._step_counter % 3) frame_data = np.array(_img)[:, :, ::-1] frame_data = face_highlight(frame_data, bboxes, 1.2) else: frame_data = frame_data.astype('int32') frame_data = np.clip(frame_data - 50, 0, 255).astype('uint8') # convert data frame into QImage h, w, c = frame_data.shape bytes_per_line = 3 * w _frame_image = QtGui.QImage(frame_data.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888) _frame_image = _frame_image.rgbSwapped() # draw frame self.scene.reset() if len(self.scene.items()) > 0: self.scene.reset_items() self.scene.addPixmap(QtGui.QPixmap.fromImage(_frame_image)) def capture(self): self.record_video.timer.stop() if self._frame_data is None: return # face detection # BGR to RGB first im = Image.fromarray(self._frame_data[:, :, ::-1]) bboxes, _ = self.detector.infer(im, min_face_size=200) if len(bboxes): faces = crop_face(self._frame_data, bboxes, 1.4) #print(lmarks) else: return if len(faces) == 0: return # draw landmarks on display image #_face_img = Image.fromarray(faces[0][:, :, ::-1]) #_, landmarks = self.detector.infer(_face_img, min_face_size=100) ##_face_img = show_bboxes(_face_img, [], landmarks) #_face_img = show_grids(_face_img, [], landmarks) #_face_img = np.array(_face_img)[:, :, ::-1] # convert data frame into QImage self._frame_data = faces[0] h, w, c = self._frame_data.shape bytes_per_line = 3 * w _frame_image = QtGui.QImage(self._frame_data.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888) #_frame_image = QtGui.QImage(_face_img.copy(), w, h, # bytes_per_line, # QtGui.QImage.Format_RGB888) _frame_image = _frame_image.rgbSwapped() image = QtGui.QPixmap.fromImage(_frame_image) mat_img = self._frame_data self.image = image.scaled(self.graphicsView.size(), QtCore.Qt.IgnoreAspectRatio) mat_img = mat_img / 127.5 - 1 self.mat_img = np.expand_dims(mat_img, axis=0) self.scene.reset() if len(self.scene.items()) > 0: self.scene.reset_items() self.scene.addPixmap(self.image) if len(self.result_scene.items()) > 0: self.result_scene.removeItem(self.result_scene.items()[-1]) self.result_scene.addPixmap(self.image) def open(self): fileName, _ = QtWidgets.QFileDialog.getOpenFileName( self, "Open File", QtCore.QDir.currentPath()) if fileName: image = QtGui.QPixmap(fileName) mat_img = cv2.imread(fileName) if image.isNull(): QtWidgets.QMessageBox.information(self, "Image Viewer", "Cannot load %s." % fileName) return # redbrush = QBrush(Qt.red) # blackpen = QPen(Qt.black) # blackpen.setWidth(5) self.image = image.scaled(self.graphicsView.size(), QtCore.Qt.IgnoreAspectRatio) mat_img = cv2.resize(mat_img, (512, 512), interpolation=cv2.INTER_CUBIC) mat_img = mat_img / 127.5 - 1 self.mat_img = np.expand_dims(mat_img, axis=0) self.scene.reset() if len(self.scene.items()) > 0: self.scene.reset_items() self.scene.addPixmap(self.image) if len(self.result_scene.items()) > 0: self.result_scene.removeItem(self.result_scene.items()[-1]) self.result_scene.addPixmap(self.image) def mask_mode(self): self.mode_select(0) def sketch_mode(self): self.mode_select(1) def stroke_mode(self): if not self.color: self.color_change_mode() self.scene.get_stk_color(self.color) self.mode_select(2) def color_change_mode(self): self.dlg.exec_() self.color = self.dlg.currentColor().name() self.pushButton_4.setStyleSheet("background-color: %s;" % self.color) self.scene.get_stk_color(self.color) def complete(self): sketch = self.make_sketch(self.scene.sketch_points) stroke = self.make_stroke(self.scene.stroke_points) mask = self.make_mask(self.scene.mask_points) if not type(self.ld_mask) == type(None): ld_mask = np.expand_dims(self.ld_mask[:, :, 0:1], axis=0) ld_mask[ld_mask > 0] = 1 ld_mask[ld_mask < 1] = 0 mask = mask + ld_mask mask[mask > 0] = 1 mask[mask < 1] = 0 mask = np.asarray(mask, dtype=np.uint8) print(mask.shape) if not type(self.ld_sk) == type(None): sketch = sketch + self.ld_sk sketch[sketch > 0] = 1 noise = self.make_noise() sketch = sketch * mask stroke = stroke * mask noise = noise * mask batch = np.concatenate([self.mat_img, sketch, stroke, mask, noise], axis=3) start_t = time.time() result = self.model.demo(self.config, batch) end_t = time.time() print('inference time : {}'.format(end_t - start_t)) result = (result + 1) * 127.5 result = np.asarray(result[0, :, :, :], dtype=np.uint8) self.output_img = result result = np.concatenate( [result[:, :, 2:3], result[:, :, 1:2], result[:, :, :1]], axis=2) qim = QtGui.QImage(result.data, result.shape[1], result.shape[0], result.strides[0], QtGui.QImage.Format_RGB888) self.result_scene.removeItem(self.result_scene.items()[-1]) self.result_scene.addPixmap(QtGui.QPixmap.fromImage(qim)) def make_noise(self): noise = np.zeros([512, 512, 1], dtype=np.uint8) noise = cv2.randn(noise, 0, 255) noise = np.asarray(noise / 255, dtype=np.uint8) noise = np.expand_dims(noise, axis=0) return noise def make_mask(self, pts): if len(pts) > 0: mask = np.zeros((512, 512, 3)) for pt in pts: cv2.line(mask, pt['prev'], pt['curr'], (255, 255, 255), 12) mask = np.asarray(mask[:, :, 0] / 255, dtype=np.uint8) mask = np.expand_dims(mask, axis=2) mask = np.expand_dims(mask, axis=0) else: mask = np.zeros((512, 512, 3)) mask = np.asarray(mask[:, :, 0] / 255, dtype=np.uint8) mask = np.expand_dims(mask, axis=2) mask = np.expand_dims(mask, axis=0) return mask def make_sketch(self, pts): if len(pts) > 0: sketch = np.zeros((512, 512, 3)) # sketch = 255*sketch for pt in pts: cv2.line(sketch, pt['prev'], pt['curr'], (255, 255, 255), 1) sketch = np.asarray(sketch[:, :, 0] / 255, dtype=np.uint8) sketch = np.expand_dims(sketch, axis=2) sketch = np.expand_dims(sketch, axis=0) else: sketch = np.zeros((512, 512, 3)) # sketch = 255*sketch sketch = np.asarray(sketch[:, :, 0] / 255, dtype=np.uint8) sketch = np.expand_dims(sketch, axis=2) sketch = np.expand_dims(sketch, axis=0) return sketch def make_stroke(self, pts): if len(pts) > 0: stroke = np.zeros((512, 512, 3)) for pt in pts: c = pt['color'].lstrip('#') color = tuple(int(c[i:i + 2], 16) for i in (0, 2, 4)) color = (color[2], color[1], color[0]) cv2.line(stroke, pt['prev'], pt['curr'], color, 4) stroke = stroke / 127.5 - 1 stroke = np.expand_dims(stroke, axis=0) else: stroke = np.zeros((512, 512, 3)) stroke = stroke / 127.5 - 1 stroke = np.expand_dims(stroke, axis=0) return stroke def arrange(self): image = np.asarray((self.mat_img[0] + 1) * 127.5, dtype=np.uint8) if len(self.scene.mask_points) > 0: for pt in self.scene.mask_points: cv2.line(image, pt['prev'], pt['curr'], (255, 255, 255), 12) if len(self.scene.stroke_points) > 0: for pt in self.scene.stroke_points: c = pt['color'].lstrip('#') color = tuple(int(c[i:i + 2], 16) for i in (0, 2, 4)) color = (color[2], color[1], color[0]) cv2.line(image, pt['prev'], pt['curr'], color, 4) if len(self.scene.sketch_points) > 0: for pt in self.scene.sketch_points: cv2.line(image, pt['prev'], pt['curr'], (0, 0, 0), 1) cv2.imwrite('tmp.jpg', image) image = QtGui.QPixmap('tmp.jpg') self.scene.history.append(3) self.scene.addPixmap(image) def save_img(self): if type(self.output_img): fileName, _ = QtWidgets.QFileDialog.getSaveFileName( self, "Save File", QtCore.QDir.currentPath()) cv2.imwrite(fileName + '.jpg', self.output_img) def undo(self): self.scene.undo() def clear(self): self.scene.reset_items() self.scene.reset() if type(self.image): self.scene.addPixmap(self.image)
class Ex(QWidget, Ui_Form): def __init__(self, model, config): super().__init__() self.setupUi(self) self.show() self.model = model self.config = config self.model.load_demo_graph(config) self.output_img = None self.mat_img = None self.ld_mask = None self.ld_sk = None self.modes = [0,0,0] self.mouse_clicked = False self.scene = GraphicsScene(self.modes) self.graphicsView.setScene(self.scene) self.graphicsView.setAlignment(Qt.AlignTop | Qt.AlignLeft) self.graphicsView.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.graphicsView.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.result_scene = QGraphicsScene() self.graphicsView_2.setScene(self.result_scene) self.graphicsView_2.setAlignment(Qt.AlignTop | Qt.AlignLeft) self.graphicsView_2.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.graphicsView_2.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.dlg = QColorDialog(self.graphicsView) self.color = None def mode_select(self, mode): for i in range(len(self.modes)): self.modes[i] = 0 self.modes[mode] = 1 def open(self): fileName, _ = QFileDialog.getOpenFileName(self, "Open File", QDir.currentPath()) if fileName: image = QPixmap(fileName) mat_img = cv2.imread(fileName) if image.isNull(): QMessageBox.information(self, "Image Viewer", "Cannot load %s." % fileName) return # redbrush = QBrush(Qt.red) # blackpen = QPen(Qt.black) # blackpen.setWidth(5) self.image = image.scaled(self.graphicsView.size(), Qt.IgnoreAspectRatio) mat_img = cv2.resize(mat_img, (512, 512), interpolation=cv2.INTER_CUBIC) mat_img = mat_img/127.5 - 1 self.mat_img = np.expand_dims(mat_img,axis=0) self.scene.reset() if len(self.scene.items())>0: self.scene.reset_items() self.scene.addPixmap(self.image) if len(self.result_scene.items())>0: self.result_scene.removeItem(self.result_scene.items()[-1]) self.result_scene.addPixmap(self.image) def mask_mode(self): self.mode_select(0) def sketch_mode(self): self.mode_select(1) def stroke_mode(self): if not self.color: self.color_change_mode() self.scene.get_stk_color(self.color) self.mode_select(2) def color_change_mode(self): self.dlg.exec_() self.color = self.dlg.currentColor().name() self.pushButton_4.setStyleSheet("background-color: %s;" % self.color) self.scene.get_stk_color(self.color) def complete(self): sketch = self.make_sketch(self.scene.sketch_points) stroke = self.make_stroke(self.scene.stroke_points) mask = self.make_mask(self.scene.mask_points) if not type(self.ld_mask)==type(None): ld_mask = np.expand_dims(self.ld_mask[:,:,0:1],axis=0) ld_mask[ld_mask>0] = 1 ld_mask[ld_mask<1] = 0 mask = mask+ld_mask mask[mask>0] = 1 mask[mask<1] = 0 mask = np.asarray(mask,dtype=np.uint8) print(mask.shape) if not type(self.ld_sk)==type(None): sketch = sketch+self.ld_sk sketch[sketch>0]=1 noise = self.make_noise() sketch = sketch*mask stroke = stroke*mask noise = noise*mask batch = np.concatenate( [self.mat_img, sketch, stroke, mask, noise],axis=3) start_t = time.time() result = self.model.demo(self.config, batch) end_t = time.time() print('inference time : {}'.format(end_t-start_t)) result = (result+1)*127.5 result = np.asarray(result[0,:,:,:],dtype=np.uint8) self.output_img = result result = np.concatenate([result[:,:,2:3],result[:,:,1:2],result[:,:,:1]],axis=2) qim = QImage(result.data, result.shape[1], result.shape[0], result.strides[0], QImage.Format_RGB888) self.result_scene.removeItem(self.result_scene.items()[-1]) self.result_scene.addPixmap(QPixmap.fromImage(qim)) def make_noise(self): noise = np.zeros([512, 512, 1],dtype=np.uint8) noise = cv2.randn(noise, 0, 255) noise = np.asarray(noise/255,dtype=np.uint8) noise = np.expand_dims(noise,axis=0) return noise def make_mask(self, pts): if len(pts)>0: mask = np.zeros((512,512,3)) for pt in pts: cv2.line(mask,pt['prev'],pt['curr'],(255,255,255),12) mask = np.asarray(mask[:,:,0]/255,dtype=np.uint8) mask = np.expand_dims(mask,axis=2) mask = np.expand_dims(mask,axis=0) else: mask = np.zeros((512,512,3)) mask = np.asarray(mask[:,:,0]/255,dtype=np.uint8) mask = np.expand_dims(mask,axis=2) mask = np.expand_dims(mask,axis=0) return mask def make_sketch(self, pts): if len(pts)>0: sketch = np.zeros((512,512,3)) # sketch = 255*sketch for pt in pts: cv2.line(sketch,pt['prev'],pt['curr'],(255,255,255),1) sketch = np.asarray(sketch[:,:,0]/255,dtype=np.uint8) sketch = np.expand_dims(sketch,axis=2) sketch = np.expand_dims(sketch,axis=0) else: sketch = np.zeros((512,512,3)) # sketch = 255*sketch sketch = np.asarray(sketch[:,:,0]/255,dtype=np.uint8) sketch = np.expand_dims(sketch,axis=2) sketch = np.expand_dims(sketch,axis=0) return sketch def make_stroke(self, pts): if len(pts)>0: stroke = np.zeros((512,512,3)) for pt in pts: c = pt['color'].lstrip('#') color = tuple(int(c[i:i+2], 16) for i in (0, 2 ,4)) color = (color[2],color[1],color[0]) cv2.line(stroke,pt['prev'],pt['curr'],color,4) stroke = stroke/127.5 - 1 stroke = np.expand_dims(stroke,axis=0) else: stroke = np.zeros((512,512,3)) stroke = stroke/127.5 - 1 stroke = np.expand_dims(stroke,axis=0) return stroke def arrange(self): image = np.asarray((self.mat_img[0]+1)*127.5,dtype=np.uint8) if len(self.scene.mask_points)>0: for pt in self.scene.mask_points: cv2.line(image,pt['prev'],pt['curr'],(255,255,255),12) if len(self.scene.stroke_points)>0: for pt in self.scene.stroke_points: c = pt['color'].lstrip('#') color = tuple(int(c[i:i+2], 16) for i in (0, 2 ,4)) color = (color[2],color[1],color[0]) cv2.line(image,pt['prev'],pt['curr'],color,4) if len(self.scene.sketch_points)>0: for pt in self.scene.sketch_points: cv2.line(image,pt['prev'],pt['curr'],(0,0,0),1) cv2.imwrite('tmp.jpg',image) image = QPixmap('tmp.jpg') self.scene.history.append(3) self.scene.addPixmap(image) def save_img(self): if type(self.output_img): fileName, _ = QFileDialog.getSaveFileName(self, "Save File", QDir.currentPath()) cv2.imwrite(fileName+'.jpg',self.output_img) def undo(self): self.scene.undo() def clear(self): self.scene.reset_items() self.scene.reset() if type(self.image): self.scene.addPixmap(self.image)
class Ex(QWidget, Ui_Form): def __init__(self, model): super().__init__() self.setupUi(self) self.show() self.model = model self.output_img = None self.mat_img = None self.ld_mask = None self.ld_sk = None self.modes = [0, 0, 0] self.mouse_clicked = False self.scene = GraphicsScene(self.modes) self.graphicsView.setScene(self.scene) self.graphicsView.setAlignment(Qt.AlignTop | Qt.AlignLeft) self.graphicsView.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.graphicsView.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.result_scene = QGraphicsScene() self.graphicsView_2.setScene(self.result_scene) self.graphicsView_2.setAlignment(Qt.AlignTop | Qt.AlignLeft) self.graphicsView_2.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.graphicsView_2.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.dlg = QColorDialog(self.graphicsView) self.color = None def mode_select(self, mode): for i in range(len(self.modes)): self.modes[i] = 0 self.modes[mode] = 1 def open(self): fileName, _ = QFileDialog.getOpenFileName(self, "Open File", QDir.currentPath()) if fileName: image = QPixmap(fileName) # print(fileName) file_label = 'test_label/' + fileName.split('/')[-1] # print(file_label) file_label = file_label.split('.')[0] + '.png' self.mat_label = cv2.imread(file_label) # print(self.mat_label) self.mat_label = cv2.cvtColor(self.mat_label, cv2.COLOR_BGR2GRAY) mat_img = cv2.imread(fileName) mat_img = cv2.cvtColor(mat_img, cv2.COLOR_BGR2RGB) if image.isNull(): QMessageBox.information(self, "Image Viewer", "Cannot load %s." % fileName) return # redbrush = QBrush(Qt.red) # blackpen = QPen(Qt.black) # blackpen.setWidth(5) self.image = image.scaled(self.graphicsView.size(), Qt.IgnoreAspectRatio) mat_img = mat_img / 127.5 - 1 self.mat_img = np.expand_dims(mat_img, axis=0) # cv2.imshow('a',self.mat_img.squeeze()) # cv2.waitKey(0) self.scene.reset() if len(self.scene.items()) > 0: self.scene.reset_items() self.scene.addPixmap(self.image) if len(self.result_scene.items()) > 0: self.result_scene.removeItem(self.result_scene.items()[-1]) self.result_scene.addPixmap(self.image) def mask_mode(self): self.mode_select(0) def sketch_mode(self): self.mode_select(1) def stroke_mode(self): if not self.color: self.color_change_mode() self.scene.get_stk_color(self.color) self.mode_select(2) def color_change_mode(self): self.dlg.exec_() self.color = self.dlg.currentColor().name() self.pushButton_4.setStyleSheet("background-color: %s;" % self.color) self.scene.get_stk_color(self.color) def complete(self): global COUNT sketch = self.make_sketch(self.scene.sketch_points) sketch = torch.FloatTensor(sketch) stroke = self.make_stroke(self.scene.stroke_points) stroke = torch.FloatTensor(stroke) mask = self.make_mask(self.scene.mask_points) mask = torch.FloatTensor(mask) self.mat_img = torch.FloatTensor(self.mat_img) if not type(self.ld_mask) == type(None): ld_mask = np.expand_dims(self.ld_mask[:, :, 0:1], axis=0) ld_mask[ld_mask > 0] = 1 ld_mask[ld_mask < 1] = 0 mask = mask + ld_mask mask[mask > 0] = 1 mask[mask < 1] = 0 mask = np.asarray(mask, dtype=np.uint8) if not type(self.ld_sk) == type(None): sketch = sketch + self.ld_sk sketch[sketch > 0] = 1 noise = self.make_noise() noise = torch.FloatTensor(noise) self.mat_label = torch.FloatTensor(self.mat_label) self.mat_label = self.mat_label.reshape(1, 512, 320, 1) start_t = time.time() result = self.model.inference( self.mat_label.permute(0, 3, 1, 2).cuda(), self.mat_img.permute(0, 3, 1, 2).cuda(), sketch.permute(0, 3, 1, 2).cuda(), stroke.permute(0, 3, 1, 2).cuda(), mask.permute(0, 3, 1, 2).cuda(), noise.permute(0, 3, 1, 2).cuda()) result = result * mask.permute( 0, 3, 1, 2).cuda() + self.mat_img.permute( 0, 3, 1, 2).cuda() * (1 - mask.permute(0, 3, 1, 2).cuda()) end_t = time.time() print('inference time : {}'.format(end_t - start_t)) result = (result + 1) * 127.5 result = result.permute(0, 2, 3, 1) result = result.cpu().numpy() result = np.asarray(result[0, :, :, :], dtype=np.uint8) self.output_img = result qim = QImage(result.data, result.shape[1], result.shape[0], result.strides[0], QImage.Format_RGB888) self.result_scene.removeItem(self.result_scene.items()[-1]) self.result_scene.addPixmap(QPixmap.fromImage(qim)) COUNT += 1 def make_noise(self): noise = np.zeros([512, 320, 1], dtype=np.uint8) noise = cv2.randn(noise, 0, 255) noise = np.asarray(noise / 255, dtype=np.uint8) noise = np.expand_dims(noise, axis=0) return noise def make_mask(self, pts): if len(pts) > 0: mask = np.zeros((512, 320, 3)) for pt in pts: cv2.line(mask, pt['prev'], pt['curr'], (255, 255, 255), 12) mask = np.asarray(mask[:, :, 0] / 255, dtype=np.uint8) mask = np.expand_dims(mask, axis=2) mask = np.expand_dims(mask, axis=0) else: mask = np.zeros((512, 320, 3)) mask = np.asarray(mask[:, :, 0] / 255, dtype=np.uint8) mask = np.expand_dims(mask, axis=2) mask = np.expand_dims(mask, axis=0) return mask def make_sketch(self, pts): if len(pts) > 0: sketch = np.zeros((512, 320, 3)) # sketch = 255*sketch for pt in pts: cv2.line(sketch, pt['prev'], pt['curr'], (255, 255, 255), 1) sketch = np.asarray(sketch[:, :, 0] / 255, dtype=np.uint8) sketch = np.expand_dims(sketch, axis=2) sketch = np.expand_dims(sketch, axis=0) else: sketch = np.zeros((512, 320, 3)) # sketch = 255*sketch sketch = np.asarray(sketch[:, :, 0] / 255, dtype=np.uint8) sketch = np.expand_dims(sketch, axis=2) sketch = np.expand_dims(sketch, axis=0) return sketch def make_stroke(self, pts): if len(pts) > 0: stroke = np.zeros((512, 320, 3)) for pt in pts: c = pt['color'].lstrip('#') color = tuple(int(c[i:i + 2], 16) for i in (0, 2, 4)) cv2.line(stroke, pt['prev'], pt['curr'], color, 4) stroke = stroke / 127.5 - 1 stroke = np.expand_dims(stroke, axis=0) else: stroke = np.zeros((512, 320, 3)) stroke = stroke / 127.5 - 1 stroke = np.expand_dims(stroke, axis=0) return stroke def arrange(self): image = np.asarray((self.mat_img[0] + 1) * 127.5, dtype=np.uint8) image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) if len(self.scene.mask_points) > 0: for pt in self.scene.mask_points: cv2.line(image, pt['prev'], pt['curr'], (255, 255, 255), 12) if len(self.scene.stroke_points) > 0: for pt in self.scene.stroke_points: c = pt['color'].lstrip('#') color = tuple(int(c[i:i + 2], 16) for i in (0, 2, 4)) color = (color[2], color[1], color[0]) cv2.line(image, pt['prev'], pt['curr'], color, 4) if len(self.scene.sketch_points) > 0: for pt in self.scene.sketch_points: cv2.line(image, pt['prev'], pt['curr'], (0, 0, 0), 1) cv2.imwrite('tmp.jpg', image) image = QPixmap('tmp.jpg') self.scene.history.append(3) self.scene.addPixmap(image) def save_img(self): if type(self.output_img): fileName, _ = QFileDialog.getSaveFileName(self, "Save File", QDir.currentPath()) cv2.imwrite(fileName + '.jpg', self.output_img) def undo(self): self.scene.undo() def clear(self): self.scene.reset_items() self.scene.reset() if type(self.image): self.scene.addPixmap(self.image)