Commit 8e8f4630 authored by dmMaze's avatar dmMaze
Browse files

fix #198

parent 26f6d2e0

modules/base.py

+6 −0
Original line number Diff line number Diff line 
from typing import Dict

from utils.logger import logger as LOGGER

import gc



GPUINTENSIVE_SET = {'cuda', 'hip'}
@@ -49,3 +50,8 @@ if hasattr(torch, 'cuda'): 
else:

    DEFAULT_DEVICE = 'cpu'



def gc_collect():

    gc.collect()

    torch.cuda.empty_cache()

    torch.cuda.ipc_collect()

modules/inpaint/__init__.py

+9 −2
Original line number Diff line number Diff line
@@ -6,7 +6,7 @@ from utils.registry import Registry 
from utils.textblock_mask import extract_ballon_mask

from utils.imgproc_utils import enlarge_window



from ..base import BaseModule, DEFAULT_DEVICE

from ..base import BaseModule, DEFAULT_DEVICE, gc_collect

from ..textdetector import TextBlock



INPAINTERS = Registry('inpainters')
@@ -44,7 +44,14 @@ class InpainterBase(BaseModule): 
                        img = img.copy()

                        img[np.where(ballon_msk > 0)] = average_bg_color

                        return img

            try:

                return self._inpaint(img, mask)

            except Exception as e:

                if isinstance(e, torch.cuda.OutOfMemoryError):

                    gc_collect()

                    return self._inpaint(img, mask)

                else:

                    raise e

        else:

            im_h, im_w = img.shape[:2]

            inpainted = np.copy(img)

ui/scenetext_manager.py

+7 −0
Original line number Diff line number Diff line
@@ -594,6 +594,13 @@ class SceneTextManager(QObject): 
        if len(selected_blks) == 0 and self.txtblkShapeControl.blk_item is not None:

            selected_blks.append(self.txtblkShapeControl.blk_item)

        if len(selected_blks) > 0:

            # img, mask = self.imgtrans_proj.img_array, self.imgtrans_proj.mask_array

            # for blk in selected_blks:

            #     x, y, w, h = blk.absBoundingRect()

            #     x2, y2 = x+w, y+h

            #     imname = str(blk.idx).zfill(2) + '.png'

            #     cv2.imwrite(imname, img[y: y2, x: x2])

            #     cv2.imwrite('mask_' + imname, mask[y: y2, x: x2])

            self.canvas.push_undo_command(DeleteBlkItemsCommand(selected_blks, mode, self))



    def onCopyBlkItems(self, pos: QPointF):