Commit 5e8bb4f7 authored by dmMaze's avatar dmMaze
Browse files

impl pen tool

parent 270d955c

ballontranslator/ui/canvas.py

+5 −0
Original line number Diff line number Diff line
@@ -355,6 +355,11 @@ class Canvas(QGraphicsScene): 
        self.inpaintLayer.setPixmap(pixmap)



    def setDrawingLayer(self, img: Union[QPixmap, np.ndarray] = None):

        

        ditems = self.drawingLayer.childItems()

        for item in ditems:

            self.removeItem(item)



        if not self.imgtrans_proj.img_valid:

            return

        if img is None:

ballontranslator/ui/drawingpanel.py

+15 −3
Original line number Diff line number Diff line
@@ -483,9 +483,8 @@ class DrawingPanel(Widget): 
            self.canvas.removeItem(stroke_item)

            return

        if self.currentTool == self.penTool:

            pass

            self.canvas.undoStack.push(StrokeItemUndoCommand(self.canvas, stroke_item))

        elif self.currentTool == self.inpaintTool:

            # stroke_item.convertToPixmapItem(remove_stroke=False, convert_mask=True)

            self.mergeInpaintStroke(stroke_item)

            if self.canvas.gv.ctrl_pressed:

                return
@@ -555,6 +554,8 @@ class DrawingPanel(Widget): 
        elif self.currentTool == self.rectTool:

            self.canvas.image_edit_mode = ImageEditMode.RectTool







    def on_finish_erasing(self, stroke_item: StrokeItem):



        # inpainted-erasing logic is essentially the same as inpainting
@@ -737,9 +738,20 @@ class DrawingPanel(Widget): 
        self.rect_inpaint_dict = None

        return super().hideEvent(e)



class StrokeItemUndoCommand(QUndoCommand):

    def __init__(self, canvas: Canvas, stroke_item: StrokeItem):

        super().__init__()

        self.stroke_pixmap = stroke_item.convertToPixmapItem()

        

    def undo(self):

        self.stroke_pixmap.hide()



    def redo(self):

        self.stroke_pixmap.show()





class InpaintUndoCommand(QUndoCommand):

    def __init__(self, canvas: Canvas, inpainted: np.ndarray, mask: np.ndarray, inpaint_rect: list):

    def __init__(self, canvas: Canvas, inpainted: np.ndarray, mask: np.ndarray, inpaint_rect: List[int]):

        super().__init__()

        self.canvas = canvas

        img_array = self.canvas.imgtrans_proj.inpainted_array

ballontranslator/ui/image_edit.py

+1 −0
Original line number Diff line number Diff line 
import numpy as np

import cv2



from qtpy.QtCore import Signal, Qt, QPointF, QSize, QPoint

from qtpy.QtWidgets import QStyleOptionGraphicsItem, QGraphicsPixmapItem, QWidget, QGraphicsPathItem, QGraphicsScene

ballontranslator/ui/mainwindow.py

+18 −4
Original line number Diff line number Diff line 
import os.path as osp

import os

import json

import cv2



from qtpy.QtWidgets import QMainWindow, QHBoxLayout, QVBoxLayout, QApplication, QStackedWidget, QWidget, QSplitter, QListWidget, QShortcut, QListWidgetItem

from qtpy.QtCore import Qt, QPoint, QSize

from qtpy.QtGui import QGuiApplication, QIcon, QCloseEvent, QKeySequence, QImage, QPainter



from utils.logger import logger as LOGGER

from .misc import ProjImgTrans

from .misc import ProjImgTrans, ndarray2pixmap, pixmap2ndarray

from .canvas import Canvas

from .configpanel import ConfigPanel

from .dl_manager import DLManager
@@ -405,13 +406,26 @@ class MainWindow(QMainWindow): 
            hide_tsc = True

            self.st_manager.txtblkShapeControl.hide()



        if not osp.exists(self.imgtrans_proj.result_dir()):

            os.makedirs(self.imgtrans_proj.result_dir())



        # save drawings to inpainted

        ditems = self.canvas.drawingLayer.childItems()

        if len(ditems) > 0:

            inpainted = self.canvas.inpaintLayer.pixmap().toImage().convertToFormat(QImage.Format.Format_ARGB32)

            painter = QPainter(inpainted)

            for ditem in ditems:

                painter.drawPixmap(ditem.pos(), ditem.pixmap())

            painter.end()

            inpainted_array = pixmap2ndarray(inpainted, keep_alpha=False)

            inpainted_path = self.imgtrans_proj.get_inpainted_path(self.imgtrans_proj.current_img)

            self.imsave_thread.saveImg(inpainted_path, inpainted_array)



        # save result: rendered text + inpainted

        painter = QPainter(img)

        painter.setRenderHint(QPainter.RenderHint.Antialiasing)

        self.canvas.render(painter)

        painter.end()

        if not osp.exists(self.imgtrans_proj.result_dir()):

            os.makedirs(self.imgtrans_proj.result_dir())



        imsave_path = self.imgtrans_proj.get_result_path(self.imgtrans_proj.current_img)

        self.imsave_thread.saveImg(imsave_path, img)

ballontranslator/ui/misc.py

+10 −4
Original line number Diff line number Diff line
@@ -26,18 +26,24 @@ def qrgb2bgr(color: Union[QColor, Tuple, List] = None) -> Tuple[int, int, int]: 
    return color



# https://stackoverflow.com/questions/45020672/convert-pyqt5-qpixmap-to-numpy-ndarray

def pixmap2ndarray(pixmap: QPixmap, keep_alpha=True):

def pixmap2ndarray(pixmap: Union[QPixmap, QImage], keep_alpha=True):

    size = pixmap.size()

    h = size.width()

    w = size.height()

    if isinstance(pixmap, QPixmap):

        qimg = pixmap.toImage().convertToFormat(QImage.Format_RGBA8888)

    else:

        qimg = pixmap



    byte_str = qimg.bits().asstring(h * w * 4)

    # byte_str.setsize(h * w * 4)

    img = np.fromstring(byte_str, dtype=np.uint8).reshape((w,h,4))

    

    if keep_alpha:

        img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)

        return img

    else:

        return np.copy(img[:, :, 0:3])

        return np.copy(img[:,:,:3])

        return cv2.cvtColor(img[:, :, 0:3], cv2.COLOR_RGB2BGR)



def ndarray2pixmap(img):

    if len(img.shape) == 2: