Skip to content

Commit

Permalink
[Feature] Implemented masked image saving functionality
Browse files Browse the repository at this point in the history
  • Loading branch information
CVHub520 committed Aug 18, 2024
1 parent 252ded0 commit f076e28
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 0 deletions.
121 changes: 121 additions & 0 deletions anylabeling/views/labeling/label_widget.py
Original file line number Diff line number Diff line change
Expand Up @@ -628,6 +628,14 @@ def __init__(
"Save cropped image. (Support rectangle/rotation/polygon shape_type)"
),
)
save_mask = action(
self.tr("&Save Masked Image"),
self.save_mask,
icon="crop",
tip=self.tr(
"Save masked image. (Support rectangle/rotation/polygon shape_type)"
),
)
label_manager = action(
self.tr("&Label Manager"),
self.label_manager,
Expand Down Expand Up @@ -1350,6 +1358,7 @@ def __init__(
overview,
None,
save_crop,
save_mask,
None,
label_manager,
None,
Expand Down Expand Up @@ -2072,6 +2081,118 @@ def save_crop(self):
error_dialog.setWindowTitle(self.tr("Error"))
error_dialog.exec_()

def save_mask(self):
if not self.filename or not self.image_list:
return

default_color = QtGui.QColor(114, 114, 114)
color = QtWidgets.QColorDialog.getColor(
default_color, self,
self.tr("Select a color to use for masking the areas")
)
if not color.isValid():
return
else:
fill_color = color.getRgb()[:3]

image_file_list, label_dir_path = self.image_list, ""
label_dir_path = osp.dirname(self.filename)
if self.output_dir:
label_dir_path = self.output_dir
save_path = osp.join(
osp.dirname(self.filename), "..", "x-anylabeling-mask-image"
)
if osp.exists(save_path):
shutil.rmtree(save_path)
os.makedirs(save_path)

progress_dialog = QProgressDialog(
self.tr("Processing..."),
self.tr("Cancel"),
0,
len(image_file_list),
)
progress_dialog.setWindowModality(Qt.WindowModal)
progress_dialog.setWindowTitle(self.tr("Progress"))
progress_dialog.setStyleSheet(
"""
QProgressDialog QProgressBar {
border: 1px solid grey;
border-radius: 5px;
text-align: center;
}
QProgressDialog QProgressBar::chunk {
background-color: orange;
}
"""
)
try:
for i, image_file in enumerate(image_file_list):
image_name = osp.basename(image_file)
label_name = osp.splitext(image_name)[0] + ".json"
label_file = osp.join(label_dir_path, label_name)
if not osp.exists(label_file):
continue
bk_image_file = osp.join(save_path, image_name)
image = cv2.imread(image_file)
with open(label_file, "r", encoding="utf-8") as f:
data = json.load(f)
shapes = data["shapes"]
for shape in shapes:
label = shape["label"]
points = shape["points"]
shape_type = shape["shape_type"]
if (label != "__mask__" or
shape_type not in ["rectangle", "polygon", "rotation"]):
continue
if (
(shape_type == "polygon" and len(points) < 3)
or (shape_type == "rotation" and len(points) != 4)
or (shape_type == "rectangle" and len(points) != 4)
):
continue
points = np.array(points, dtype=np.int32)
if shape['shape_type'] == 'rectangle':
top_left = tuple(points[0])
bottom_right = tuple(points[2])
cv2.rectangle(image, top_left, bottom_right, fill_color, thickness=-1)
elif shape['shape_type'] == 'rotation':
rect = cv2.minAreaRect(points)
box = cv2.boxPoints(rect)
box = np.int0(box)
cv2.drawContours(image, [box], 0, fill_color, thickness=cv2.FILLED)
elif shape['shape_type'] == 'polygon':
cv2.fillPoly(image, [points], fill_color)
cv2.imwrite(bk_image_file, image)
# Update progress bar
progress_dialog.setValue(i)
if progress_dialog.wasCanceled():
break

# Hide the progress dialog after processing is done
progress_dialog.close()
# Show success message
save_path = osp.realpath(save_path)
msg_box = QMessageBox()
msg_box.setIcon(QMessageBox.Information)
msg_box.setText(self.tr("Masking completed successfully!"))
msg_box.setInformativeText(
self.tr(f"Results have been saved to:\n{save_path}")
)
msg_box.setWindowTitle(self.tr("Success"))
msg_box.exec_()

except Exception as e:
progress_dialog.close()
error_dialog = QMessageBox()
error_dialog.setIcon(QMessageBox.Critical)
error_dialog.setText(
self.tr("Error occurred while masking image.")
)
error_dialog.setInformativeText(str(e))
error_dialog.setWindowTitle(self.tr("Error"))
error_dialog.exec_()

def label_manager(self):
modify_label_dialog = LabelModifyDialog(
parent=self, opacity=LABEL_OPACITY
Expand Down
7 changes: 7 additions & 0 deletions docs/en/user_guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
* [5.2 Save Sub-Images](#52-save-sub-images)
* [5.3 Object Manager](#53-object-manager)
* [5.4 Shape Type Conversion](#54-shape-type-conversion)
* [5.5 Saving Masked Images](#55-saving-masked-images)
* [6. Help and Language](#6-help-and-language)
* [6.1 Getting Version Information](#61-getting-version-information)
* [6.2 Setting Software Language](#62-setting-software-language)
Expand Down Expand Up @@ -506,6 +507,12 @@ These actions can be accessed through the **Tools** -> **Object Manager** option

Note: The `Rotated Box to Rectangular Box` and `Polygon to Bounding Box` conversions use the maximum bounding rectangle by default, which may result in the loss of some label information, such as rotation angle. These actions are **irreversible**, so use them with caution.

### 5.5 Saving Masked Images

To implement the masked image saving feature, follow the steps below:
1. For areas that require `masking`, they can be identified using `rectangle`, `rotation`, or `polygon` objects, with the label name set to `__mask__`;
2. Click on the **Tools** -> **Save Masked Image** option in the top menu bar to set the fill color block. The default value is (114, 114, 114). The resulting images are saved by default in the `x-anylabeling-mask-image` folder.

## 6. Help and Language

### 6.1 Getting Version Information
Expand Down
7 changes: 7 additions & 0 deletions docs/zh_cn/user_guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
* [5.2 保存子图](#52-保存子图)
* [5.3 对象管理器](#53-对象管理器)
* [5.4 标签类型转换](#54-标签类型转换)
* [5.5 保存遮罩图像](#55-保存遮罩图像)
* [6. 帮助及语言](#6-帮助及语言)
* [6.1 获取版本信息](#61-获取版本信息)
* [6.2 设置软件语言](#62-设置软件语言)
Expand Down Expand Up @@ -502,6 +503,12 @@ imext = .jpg # 图片文件扩展名

注意:`旋转框转矩形框``多边形框转多边形框` 操作默认取最大外接矩,因此会丢失一些标签信息,如旋转角度等。这些操作为**不可逆**,请谨慎使用。

### 5.5 保存遮罩图像

实现遮罩图像存储功能,请按照以下步骤操作:

1. 对于需执行 `mask` 操作的区域,使用 `rectangle``rotation``polygon` 对象进行标记,并将标签命名为 `__mask__`
2. 在顶部菜单栏选择 **工具** -> **保存遮罩图像**,然后设定填充色块(默认为 (114, 114, 114))。生成的图像将默认保存在 `x-anylabeling-mask-image` 文件夹中。

## 6. 帮助及语言

Expand Down

0 comments on commit f076e28

Please sign in to comment.